Improve DataSource TS

packages/core/src/data_sources/model/DataRecords.ts

@@ -3,7 +3,7 @@ import { DataRecordProps } from '../types';
 import DataRecord from './DataRecord';
 import DataSource from './DataSource';
 
-export default class DataRecords extends Collection<DataRecord> {
+export default class DataRecords<T extends DataRecordProps = DataRecordProps> extends Collection<DataRecord<T>> {
   dataSource: DataSource;
 
   constructor(models: DataRecord[] | DataRecordProps[], options: { dataSource: DataSource }) {

packages/core/src/data_sources/model/DataSource.ts

@@ -31,14 +31,18 @@
 
 import { AddOptions, collectionEvents, CombinedModelConstructorOptions, Model, RemoveOptions } from '../../common';
 import EditorModel from '../../editor/model/Editor';
-import { DataRecordProps, DataSourceProps, DataSourceTransformers } from '../types';
+import { DataSourceProps } from '../types';
+import { DataSourceTransformers, DataSourceType, SingleRecordType } from '../types';
 import DataRecord from './DataRecord';
 import DataRecords from './DataRecords';
 import DataSources from './DataSources';
 
 interface DataSourceOptions extends CombinedModelConstructorOptions<{ em: EditorModel }, DataSource> {}
 
-export default class DataSource extends Model<DataSourceProps> {
+export default class DataSource<
+  DS extends DataSourceType = DataSourceType,
+  DR extends SingleRecordType<DS['records']> = SingleRecordType<DS['records']>,
+> extends Model<DS> {
   transformers: DataSourceTransformers;
 
   /**
@@ -52,7 +56,7 @@ export default class DataSource extends Model<DataSourceProps> {
     return {
       records: [],
       transformers: {},
-    };
+    } as unknown as Partial<DS>;
   }
 
   /**
@@ -64,13 +68,19 @@ export default class DataSource extends Model<DataSourceProps> {
    * @param {DataSourceOptions} opts - Options to initialize the data source.
    * @name constructor
    */
-  constructor(props: DataSourceProps, opts: DataSourceOptions) {
-    super(props, opts);
+  constructor(props: DataSourceProps<DS>, opts: DataSourceOptions) {
+    super(
+      {
+        ...props,
+        records: [],
+      } as unknown as DS,
+      opts,
+    );
     const { records, transformers } = props;
     this.transformers = transformers || {};
 
     if (!(records instanceof DataRecords)) {
-      this.set({ records: new DataRecords(records!, { dataSource: this }) });
+      this.set({ records: new DataRecords(records!, { dataSource: this }) } as Partial<DS>);
     }
 
     this.listenTo(this.records, 'add', this.onAdd);
@@ -84,7 +94,7 @@ export default class DataSource extends Model<DataSourceProps> {
    * @name records
    */
   get records() {
-    return this.attributes.records as DataRecords;
+    return this.attributes.records as NonNullable<DS['records']>;
   }
 
   /**
@@ -117,7 +127,7 @@ export default class DataSource extends Model<DataSourceProps> {
    * @returns {DataRecord} The added data record.
    * @name addRecord
    */
-  addRecord(record: DataRecordProps, opts?: AddOptions) {
+  addRecord(record: DR, opts?: AddOptions) {
     return this.records.add(record, opts);
   }
 
@@ -128,9 +138,8 @@ export default class DataSource extends Model<DataSourceProps> {
    * @returns {DataRecord | undefined} The data record, or `undefined` if no record is found with the given ID.
    * @name getRecord
    */
-  getRecord(id: string | number): DataRecord | undefined {
-    const record = this.records.get(id);
-    return record;
+  getRecord(id: string | number) {
+    return this.records.get(id) as DR | undefined;
   }
 
   /**
@@ -141,7 +150,7 @@ export default class DataSource extends Model<DataSourceProps> {
    * @name getRecords
    */
   getRecords() {
-    return [...this.records.models].map((record) => this.getRecord(record.id));
+    return [...this.records.models].map((record) => this.getRecord(record.id)!);
   }
 
   /**
@@ -168,7 +177,7 @@ export default class DataSource extends Model<DataSourceProps> {
    * @returns {Array<DataRecord>} An array of the added data records.
    * @name setRecords
    */
-  setRecords(records: Array<DataRecordProps>) {
+  setRecords(records: DR[]) {
     this.records.reset([], { silent: true });
 
     records.forEach((record) => {

packages/core/src/data_sources/types.ts

@@ -1,4 +1,4 @@
-import { ObjectAny } from '../common';
+import { Collection, ObjectAny } from '../common';
 import ComponentDataVariable from './model/ComponentDataVariable';
 import DataRecord from './model/DataRecord';
 import DataRecords from './model/DataRecords';
@@ -24,17 +24,12 @@ export interface DataVariableListener {
   event: string;
 }
 
-export interface DataSourceProps {
+interface BaseDataSource {
   /**
    * DataSource id.
    */
   id: string;
 
-  /**
-   * DataSource records.
-   */
-  records?: DataRecords | DataRecord[] | DataRecordProps[];
-
   /**
    * DataSource validation and transformation factories.
    */
@@ -45,6 +40,14 @@ export interface DataSourceProps {
    */
   skipFromStorage?: boolean;
 }
+export interface DataSourceType<DR extends DataRecordProps = DataRecordProps> extends BaseDataSource {
+  records: DataRecords<DR>;
+}
+export interface DataSourceProps<DS extends DataSourceType = DataSourceType> extends BaseDataSource {
+  records?: DataRecords<ExtractRecordType<DS>> | DataRecord<ExtractRecordType<DS>>[] | ExtractRecordType<DS>[];
+}
+export type ExtractRecordType<T> = T extends { records: DataRecords<infer DR> } ? DR : never;
+export type SingleRecordType<T> = T extends Collection<infer U> ? U : never;
 
 export interface DataSourceTransformers {
   onRecordSetValue?: (args: { id: string | number; key: string; value: any }) => any;