Merge pull request #1678 from westito/custom_parent_class

Add custom parent class option to generated data classes

.gitignore

@@ -13,3 +13,4 @@ benchmark_results.json
 flutter_export_environment.sh
 # Local Netlify folder
 .netlify
+.DS_Store

drift/CHANGELOG.md

@@ -1,3 +1,8 @@
+## 1.5.0-dev
+
+- Add `DataClassName.extending` to control the superclass of generated row
+  classes.
+
 ## 1.4.0
 
 - Most methods to compose statements are now available as an extension on

drift/lib/src/dsl/table.dart

@@ -184,9 +184,42 @@ class DataClassName {
   /// {@macro drift_custom_data_class}
   final String name;
 
+  /// The parent type of the data class generated by drift.
+  ///
+  /// The [extending] type must refer to an interface type (usually just a
+  /// class name), and the parent class must extend [DataClass].
+  ///
+  /// The extended class can optionally have a type parameter, which is
+  /// instantiated to the actual data class generated by drift.
+  ///
+  /// For example,
+  ///
+  /// ```dart
+  ///  abstract class BaseModel extends DataClass {
+  ///    abstract final String id;
+  ///  }
+  ///
+  ///  abstract class TypedBaseModel<T> extends DataClass {
+  ///
+  ///  }
+  ///
+  ///  @DataClassName('Company', extending: BaseModel)
+  ///  class Companies extends Table {
+  ///    TextColumn get id => text()();
+  ///    TextColumn get name => text().named('name')();
+  ///  }
+  ///
+  ///  // The actual generated class will extend `TypedBaseModel<Employee>`.
+  ///  @DataClassName('Employee', extending: TypedBaseModel)
+  ///  class Employees extends Table {
+  ///    TextColumn get id => text()();
+  ///  }
+  /// ```
+  final Type? extending;
+
   /// Customize the data class name for a given table.
   /// {@macro drift_custom_data_class}
-  const DataClassName(this.name);
+  const DataClassName(this.name, {this.extending});
 }
 
 /// An annotation specifying an existing class to be used as a data class.
@@ -239,6 +272,9 @@ class DriftView {
   /// {@macro drift_custom_data_class}
   final String? dataClassName;
 
+  /// The parent class of generated data class. Class must extends [DataClass]!
+  final Type? extending;
+
   /// Customize view name and data class name
-  const DriftView({this.name, this.dataClassName});
+  const DriftView({this.name, this.dataClassName, this.extending});
 }

drift/pubspec.yaml

@@ -1,6 +1,6 @@
 name: drift
 description: Drift is a reactive library to store relational data in Dart and Flutter applications.
-version: 1.4.0
+version: 1.5.0-dev
 repository: https://github.com/simolus3/moor
 homepage: https://drift.simonbinder.eu/
 issue_tracker: https://github.com/simolus3/moor/issues

drift_dev/lib/src/analyzer/dart/parser.dart

@@ -5,10 +5,10 @@ import 'package:analyzer/dart/element/nullability_suffix.dart';
 import 'package:analyzer/dart/element/type.dart';
 import 'package:collection/collection.dart';
 import 'package:drift_dev/moor_generator.dart';
+import 'package:drift_dev/src/analyzer/data_class.dart';
 import 'package:drift_dev/src/analyzer/errors.dart';
 import 'package:drift_dev/src/analyzer/runner/steps.dart';
 import 'package:drift_dev/src/utils/exception.dart';
-import 'package:drift_dev/src/utils/names.dart';
 import 'package:drift_dev/src/utils/type_utils.dart';
 import 'package:meta/meta.dart';
 import 'package:recase/recase.dart';
@@ -19,9 +19,9 @@ import '../custom_row_class.dart';
 
 part 'column_parser.dart';
 part 'table_parser.dart';
-part 'view_parser.dart';
 part 'use_dao_parser.dart';
 part 'use_moor_parser.dart';
+part 'view_parser.dart';
 
 class MoorDartParser {
   final ParseDartStep step;

drift_dev/lib/src/analyzer/dart/table_parser.dart

@@ -21,6 +21,7 @@ class TableParser {
       sqlName: sqlName,
       dartTypeName: dataClassInfo.enforcedName,
       existingRowClass: dataClassInfo.existingClass,
+      customParentClass: dataClassInfo.extending,
       primaryKey: primaryKey,
       overrideWithoutRowId: await _overrideWithoutRowId(element),
       declaration: DartTableDeclaration(element, base.step.file),
@@ -68,12 +69,15 @@ class TableParser {
     }
 
     String name;
+    String? customParentClass;
     FoundDartClass? existingClass;
     String? constructorInExistingClass;
     bool? generateInsertable;
 
     if (dataClassName != null) {
       name = dataClassName.getField('name')!.toStringValue()!;
+      customParentClass =
+          parseCustomParentClass(name, dataClassName, element, base);
     } else {
       name = dataClassNameForClassName(element.name);
     }
@@ -100,7 +104,7 @@ class TableParser {
         ? null
         : validateExistingClass(columns, existingClass,
             constructorInExistingClass!, generateInsertable!, base.step);
-    return _DataClassInformation(name, verified);
+    return _DataClassInformation(name, customParentClass, verified);
   }
 
   Future<String?> _parseTableName(ClassElement element) async {
@@ -238,9 +242,14 @@ class TableParser {
 
 class _DataClassInformation {
   final String enforcedName;
+  final String? extending;
   final ExistingRowClass? existingClass;
 
-  _DataClassInformation(this.enforcedName, this.existingClass);
+  _DataClassInformation(
+    this.enforcedName,
+    this.extending,
+    this.existingClass,
+  );
 }
 
 extension on Element {

drift_dev/lib/src/analyzer/dart/view_parser.dart

@@ -20,6 +20,7 @@ class ViewParser {
       name: name,
       dartTypeName: dataClassInfo.enforcedName,
       existingRowClass: dataClassInfo.existingClass,
+      customParentClass: dataClassInfo.extending,
       entityInfoName: '\$${element.name}View',
       viewQuery: query,
     );
@@ -35,20 +36,21 @@ class ViewParser {
   _DataClassInformation _readDataClassInformation(
       List<MoorColumn> columns, ClassElement element) {
     DartObject? useRowClass;
-    String? dataClassName;
+    DartObject? driftView;
+    String? customParentClass;
 
     for (final annotation in element.metadata) {
       final computed = annotation.computeConstantValue();
       final annotationClass = computed!.type!.element!.name;
 
       if (annotationClass == 'DriftView') {
-        dataClassName = computed.getField('dataClassName')?.toStringValue();
+        driftView = computed;
       } else if (annotationClass == 'UseRowClass') {
         useRowClass = computed;
       }
     }
 
-    if (dataClassName != null && useRowClass != null) {
+    if (driftView != null && useRowClass != null) {
       base.step.reportError(ErrorInDartCode(
         message: "A table can't be annotated with both @DataClassName and "
             '@UseRowClass',
@@ -60,7 +62,15 @@ class ViewParser {
     String? constructorInExistingClass;
     bool? generateInsertable;
 
-    var name = dataClassName ?? dataClassNameForClassName(element.name);
+    var name = dataClassNameForClassName(element.name);
+
+    if (driftView != null) {
+      final dataClassName =
+          driftView.getField('dataClassName')?.toStringValue();
+      name = dataClassName ?? dataClassNameForClassName(element.name);
+      customParentClass =
+          parseCustomParentClass(name, driftView, element, base);
+    }
 
     if (useRowClass != null) {
       final type = useRowClass.getField('type')!.toTypeValue();
@@ -84,7 +94,7 @@ class ViewParser {
         ? null
         : validateExistingClass(columns, existingClass,
             constructorInExistingClass!, generateInsertable!, base.step);
-    return _DataClassInformation(name, verified);
+    return _DataClassInformation(name, customParentClass, verified);
   }
 
   Future<String> _parseViewName(ClassElement element) async {

drift_dev/lib/src/analyzer/data_class.dart

@@ -0,0 +1,85 @@
+import 'package:analyzer/dart/constant/value.dart';
+import 'package:analyzer/dart/element/element.dart';
+import 'package:analyzer/dart/element/type.dart';
+import 'package:drift_dev/src/analyzer/dart/parser.dart';
+import 'package:drift_dev/src/analyzer/errors.dart';
+import 'package:drift_dev/src/utils/type_utils.dart';
+
+String dataClassNameForClassName(String tableName) {
+  // This implementation is very primitive at the moment. The basic idea is
+  // that, very often, table names are formed from the plural of the entity
+  // they're storing (users, products, ...). We try to find the singular word
+  // from the table name.
+
+  // todo we might want to implement some edge cases according to
+  // https://en.wikipedia.org/wiki/English_plurals
+
+  if (tableName.endsWith('s')) {
+    return tableName.substring(0, tableName.length - 1);
+  }
+
+  // Default behavior if the table name is not a valid plural.
+  return '${tableName}Data';
+}
+
+String? parseCustomParentClass(String dartTypeName, DartObject dataClassName,
+    ClassElement element, MoorDartParser base) {
+  final extending = dataClassName.getField('extending');
+  if (extending != null && !extending.isNull) {
+    final extendingType = extending.toTypeValue();
+    if (extendingType is InterfaceType) {
+      final superType = extendingType.allSupertypes
+          .any((type) => isFromMoor(type) && type.element.name == 'DataClass');
+      if (!superType) {
+        base.step.reportError(
+          ErrorInDartCode(
+            message: 'Parameter `extending` in @DataClassName must be subtype '
+                'of DataClass',
+            affectedElement: element,
+          ),
+        );
+        return null;
+      }
+
+      if (extendingType.typeArguments.length > 1) {
+        base.step.reportError(
+          ErrorInDartCode(
+            message: 'Parameter `extending` in @DataClassName must have zero or'
+                ' one type parameter',
+            affectedElement: element,
+          ),
+        );
+        return null;
+      }
+
+      final className = extendingType.element.name;
+      if (extendingType.typeArguments.length == 1) {
+        final genericType = extendingType.typeArguments[0].element?.name;
+        if (genericType == 'Object' || genericType == 'dynamic') {
+          return '$className<$dartTypeName>';
+        } else {
+          base.step.reportError(
+            ErrorInDartCode(
+              message: 'Parameter `extending` in @DataClassName can only be '
+                  'provided as `$className<Object>`, `$className<dynamic>` or '
+                  'without declared type parameter (`$className`)',
+              affectedElement: element,
+            ),
+          );
+          return null;
+        }
+      }
+
+      return className;
+    } else {
+      base.step.reportError(
+        ErrorInDartCode(
+          message: 'Parameter `extending` in @DataClassName must be used with '
+              'a class',
+          affectedElement: element,
+        ),
+      );
+    }
+  }
+  return null;
+}

drift_dev/lib/src/analyzer/moor/create_table_reader.dart

@@ -1,11 +1,11 @@
 import 'package:analyzer/dart/element/nullability_suffix.dart';
 import 'package:analyzer/dart/element/type.dart';
 import 'package:drift_dev/moor_generator.dart';
+import 'package:drift_dev/src/analyzer/data_class.dart';
 import 'package:drift_dev/src/analyzer/errors.dart';
 import 'package:drift_dev/src/analyzer/runner/steps.dart';
 import 'package:drift_dev/src/analyzer/sql_queries/type_mapping.dart';
 import 'package:drift_dev/src/backends/backend.dart';
-import 'package:drift_dev/src/utils/names.dart';
 import 'package:drift_dev/src/utils/string_escaper.dart';
 import 'package:drift_dev/src/utils/type_converter_hint.dart';
 import 'package:drift_dev/src/utils/type_utils.dart';

drift_dev/lib/src/model/base_entity.dart

@@ -49,6 +49,9 @@ abstract class MoorEntityWithResultSet extends MoorSchemaEntity {
   /// The existing class designed to hold a row, if there is any.
   ExistingRowClass? get existingRowClass;
 
+  /// Class that added to data class as implementation
+  String? get customParentClass;
+
   /// The name of the Dart class storing the right column getters for this type.
   ///
   /// This class is equal to, or a superclass of, [entityInfoName].

drift_dev/lib/src/model/table.dart

@@ -27,6 +27,9 @@ class MoorTable extends MoorEntityWithResultSet {
   @override
   final ExistingRowClass? existingRowClass;
 
+  @override
+  final String? customParentClass;
+
   /// If [fromClass] is null, another source to use when determining the name
   /// of this table in generated Dart code.
   final String? _overriddenName;
@@ -149,6 +152,7 @@ class MoorTable extends MoorEntityWithResultSet {
     this.overrideDontWriteConstraints,
     this.declaration,
     this.existingRowClass,
+    this.customParentClass,
     this.isStrict = false,
   }) : _overriddenName = overriddenName {
     _attachToConverters();

drift_dev/lib/src/model/view.dart

@@ -1,7 +1,7 @@
+import 'package:drift_dev/src/analyzer/data_class.dart';
 import 'package:drift_dev/src/analyzer/options.dart';
 import 'package:drift_dev/src/analyzer/runner/file_graph.dart';
 import 'package:drift_dev/src/analyzer/runner/results.dart';
-import 'package:drift_dev/src/utils/names.dart';
 import 'package:recase/recase.dart';
 import 'package:sqlparser/sqlparser.dart';
 
@@ -36,6 +36,9 @@ class MoorView extends MoorEntityWithResultSet {
   @override
   ExistingRowClass? existingRowClass;
 
+  @override
+  final String? customParentClass;
+
   final ViewQueryInformation? viewQuery;
 
   MoorView({
@@ -44,6 +47,7 @@ class MoorView extends MoorEntityWithResultSet {
     required this.dartTypeName,
     required this.entityInfoName,
     this.existingRowClass,
+    this.customParentClass,
     this.viewQuery,
   });
 

drift_dev/lib/src/utils/names.dart

@@ -1,16 +0,0 @@
-String dataClassNameForClassName(String tableName) {
-  // This implementation is very primitive at the moment. The basic idea is
-  // that, very often, table names are formed from the plural of the entity
-  // they're storing (users, products, ...). We try to find the singular word
-  // from the table name.
-
-  // todo we might want to implement some edge cases according to
-  // https://en.wikipedia.org/wiki/English_plurals
-
-  if (tableName.endsWith('s')) {
-    return tableName.substring(0, tableName.length - 1);
-  }
-
-  // Default behavior if the table name is not a valid plural.
-  return '${tableName}Data';
-}

drift_dev/lib/src/writer/tables/data_class_writer.dart

@@ -24,7 +24,9 @@ class DataClassWriter {
       : 'driftRuntimeOptions';
 
   void write() {
-    _buffer.write('class ${table.dartTypeName} extends DataClass ');
+    final parentClass = table.customParentClass ?? 'DataClass';
+    _buffer.write('class ${table.dartTypeName} extends $parentClass ');
+
     if (isInsertable) {
       // The data class is only an insertable if we can actually insert rows
       // into the target entity.

drift_dev/pubspec.yaml

@@ -1,6 +1,6 @@
 name: drift_dev
 description: Dev-dependency for users of drift. Contains a the generator and development tools.
-version: 1.4.0
+version: 1.5.0-dev
 repository: https://github.com/simolus3/moor
 homepage: https://drift.simonbinder.eu/
 issue_tracker: https://github.com/simolus3/moor/issues
@@ -25,7 +25,7 @@ dependencies:
   io: ^1.0.3
 
   # Drift-specific analysis and apis
-  drift: '>=1.4.0 <1.5.0'
+  drift: '>=1.5.0 <1.6.0'
   sqlite3: '>=0.1.6 <2.0.0'
   sqlparser: ^0.20.0
 

drift_dev/test/analyzer/dart/custom_row_classes_test.dart

@@ -148,6 +148,98 @@ class Cls extends HasBar {
 
   Cls(this.foo, int bar): super(bar);
 }
+''',
+      'a|lib/custom_parent_class_no_error.dart': '''
+import 'package:drift/drift.dart';
+  
+abstract class BaseModel extends DataClass {
+  abstract final String id;
+}
+
+@DataClassName('Company', extending: BaseModel)
+class Companies extends Table {
+  TextColumn get id => text()();
+  TextColumn get name => text().named('name')();
+}
+''',
+      'a|lib/custom_parent_class_typed_no_error.dart': '''
+import 'package:drift/drift.dart';
+  
+abstract class BaseModel<T> extends DataClass {
+  abstract final String id;
+}
+
+@DataClassName('Company', extending: BaseModel)
+class Companies extends Table {
+  TextColumn get id => text()();
+  TextColumn get name => text().named('name')();
+}
+''',
+      'a|lib/custom_parent_class_no_super.dart': '''
+import 'package:drift/drift.dart';
+  
+abstract class BaseModel {
+  abstract final String id;
+}
+
+@DataClassName('Company', extending: BaseModel)
+class Companies extends Table {
+  TextColumn get id => text()();
+  TextColumn get name => text().named('name')();
+}
+''',
+      'a|lib/custom_parent_class_wrong_super.dart': '''
+import 'package:drift/drift.dart';
+  
+class Test {
+}
+
+abstract class BaseModel extends Test {
+  abstract final String id;
+}
+
+@DataClassName('Company', extending: BaseModel)
+class Companies extends Table {
+  TextColumn get id => text()();
+  TextColumn get name => text().named('name')();
+}
+''',
+      'a|lib/custom_parent_class_typed_wrong_type_arg.dart': '''
+import 'package:drift/drift.dart';
+  
+abstract class BaseModel<T> extends DataClass {
+  abstract final String id;
+}
+
+@DataClassName('Company', extending: BaseModel<String>)
+class Companies extends Table {
+  TextColumn get id => text()();
+  TextColumn get name => text().named('name')();
+}
+''',
+      'a|lib/custom_parent_class_two_type_argument.dart': '''
+import 'package:drift/drift.dart';
+  
+abstract class BaseModel<T, D> extends DataClass {
+  abstract final String id;
+}
+
+@DataClassName('Company', extending: BaseModel)
+class Companies extends Table {
+  TextColumn get id => text()();
+  TextColumn get name => text().named('name')();
+}
+''',
+      'a|lib/custom_parent_class_not_class.dart': '''
+import 'package:drift/drift.dart';
+
+typedef NotClass = void Function();
+  
+@DataClassName('Company', extending: NotClass)
+class Companies extends Table {
+  TextColumn get id => text()();
+  TextColumn get name => text().named('name')();
+}
 ''',
     });
   });
@@ -268,4 +360,88 @@ class Cls extends HasBar {
     final file = await state.analyze('package:a/insertable_valid.dart');
     expect(file.errors.errors, isEmpty);
   });
+
+  group('custom data class parent', () {
+    test('check valid', () async {
+      final file =
+          await state.analyze('package:a/custom_parent_class_no_error.dart');
+      expect(file.errors.errors, isEmpty);
+    });
+
+    test('check valid with type argument', () async {
+      final file = await state
+          .analyze('package:a/custom_parent_class_typed_no_error.dart');
+      expect(file.errors.errors, isEmpty);
+    });
+
+    test('check extends DataClass (no super)', () async {
+      final file =
+          await state.analyze('package:a/custom_parent_class_no_super.dart');
+
+      expect(
+        file.errors.errors,
+        contains(isA<ErrorInDartCode>().having(
+            (e) => e.message,
+            'message',
+            contains('Parameter `extending` in '
+                '@DataClassName must be subtype of DataClass'))),
+      );
+    });
+
+    test('extends DataClass (wrong super)', () async {
+      final file =
+          await state.analyze('package:a/custom_parent_class_wrong_super.dart');
+
+      expect(
+        file.errors.errors,
+        contains(isA<ErrorInDartCode>().having(
+            (e) => e.message,
+            'message',
+            contains('Parameter `extending` in '
+                '@DataClassName must be subtype of DataClass'))),
+      );
+    });
+
+    test('wrong type argument in extending', () async {
+      final file = await state
+          .analyze('package:a/custom_parent_class_typed_wrong_type_arg.dart');
+
+      expect(
+        file.errors.errors,
+        contains(isA<ErrorInDartCode>().having(
+            (e) => e.message,
+            'message',
+            contains('Parameter `extending` in @DataClassName can only be '
+                'provided as'))),
+      );
+    });
+
+    test('two type arguments in parent class', () async {
+      final file = await state
+          .analyze('package:a/custom_parent_class_two_type_argument.dart');
+
+      expect(
+        file.errors.errors,
+        contains(isA<ErrorInDartCode>().having(
+            (e) => e.message,
+            'message',
+            contains('Parameter `extending` in @DataClassName must have zero '
+                'or one type parameter'))),
+      );
+    });
+
+    test('not a class in extending', () async {
+      final file =
+          await state.analyze('package:a/custom_parent_class_not_class.dart');
+
+      expect(
+        file.errors.errors,
+        contains(isA<ErrorInDartCode>().having(
+            (e) => e.message,
+            'message',
+            contains('Parameter `extending` in @DataClassName must be used '
+                'with a class'))),
+      );
+    });
+  });
 }