From 1ebf65157a7f4f2db43572060f8ed8dfa2e10921 Mon Sep 17 00:00:00 2001
From: Artur Arseniev <artur.catch@hotmail.it>
Date: Tue, 5 Nov 2024 18:03:17 +0400
Subject: [PATCH] Datasource types (#6295)

* Add datasource types

* Refactor types

* Remove unused type

---------

Co-authored-by: mohamedsalem401 <salemmohamed780@gmail.com>
---
 .../core/src/data_sources/model/DataSource.ts | 32 ++++++++++++-------
 packages/core/src/data_sources/types.ts       | 20 ++++++------
 2 files changed, 31 insertions(+), 21 deletions(-)

diff --git a/packages/core/src/data_sources/model/DataSource.ts b/packages/core/src/data_sources/model/DataSource.ts
index a585558e9b..618ee30688 100644
--- a/packages/core/src/data_sources/model/DataSource.ts
+++ b/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<DS extends DataSourceProps = DataSourceProps> extends Model<DS> {
+export default class DataSource<
+  DS extends DataSourceType = DataSourceType,
+  DR extends SingleRecordType<DS['records']> = SingleRecordType<DS['records']>,
+> extends Model<DS> {
   transformers: DataSourceTransformers;
 
   /**
@@ -64,8 +68,14 @@ export default class DataSource<DS extends DataSourceProps = DataSourceProps> ex
    * @param {DataSourceOptions} opts - Options to initialize the data source.
    * @name constructor
    */
-  constructor(props: DS, 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 || {};
 
@@ -84,7 +94,7 @@ export default class DataSource<DS extends DataSourceProps = DataSourceProps> ex
    * @name records
    */
   get records() {
-    return this.attributes.records as NonNullable<DS['_records']>;
+    return this.attributes.records as NonNullable<DS['records']>;
   }
 
   /**
@@ -117,7 +127,7 @@ export default class DataSource<DS extends DataSourceProps = DataSourceProps> ex
    * @returns {DataRecord} The added data record.
    * @name addRecord
    */
-  addRecord<T extends DataRecordProps = DataRecordProps>(record: T, opts?: AddOptions) {
+  addRecord(record: DR, opts?: AddOptions) {
     return this.records.add(record, opts);
   }
 
@@ -128,8 +138,8 @@ export default class DataSource<DS extends DataSourceProps = DataSourceProps> ex
    * @returns {DataRecord | undefined} The data record, or `undefined` if no record is found with the given ID.
    * @name getRecord
    */
-  getRecord<T extends DataRecordProps = DataRecordProps>(id: string | number) {
-    return this.records.get(id) as DataRecord<T> | undefined;
+  getRecord(id: string | number) {
+    return this.records.get(id) as DR | undefined;
   }
 
   /**
@@ -139,8 +149,8 @@ export default class DataSource<DS extends DataSourceProps = DataSourceProps> ex
    * @returns {Array<DataRecord | undefined>} An array of data records.
    * @name getRecords
    */
-  getRecords<T extends DataRecordProps = DataRecordProps>() {
-    return [...this.records.models].map((record) => this.getRecord<T>(record.id)!);
+  getRecords() {
+    return [...this.records.models].map((record) => this.getRecord(record.id)!);
   }
 
   /**
@@ -167,7 +177,7 @@ export default class DataSource<DS extends DataSourceProps = DataSourceProps> ex
    * @returns {Array<DataRecord>} An array of the added data records.
    * @name setRecords
    */
-  setRecords<T extends DataRecordProps = DataRecordProps>(records: T[]) {
+  setRecords(records: DR[]) {
     this.records.reset([], { silent: true });
 
     records.forEach((record) => {
diff --git a/packages/core/src/data_sources/types.ts b/packages/core/src/data_sources/types.ts
index 57f358adea..ec9508c7a9 100644
--- a/packages/core/src/data_sources/types.ts
+++ b/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<DR extends DataRecordProps = DataRecordProps> {
+interface BaseDataSource {
   /**
    * DataSource id.
    */
   id: string;
 
-  /**
-   * DataSource records.
-   */
-  records?: DataRecords<DR> | DataRecord<DR>[] | DR[];
-
   /**
    * DataSource validation and transformation factories.
    */
@@ -44,10 +39,15 @@ export interface DataSourceProps<DR extends DataRecordProps = DataRecordProps> {
    * If true will store the data source in the GrapesJS project.json file.
    */
   skipFromStorage?: boolean;
-
-  // for TS
-  _records?: DataRecords<DR>;
 }
+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;