Merge pull request #2913 from ValentinVignal/drift-dev/name-the-companion-class

feat: Allow to name companion classes

drift/CHANGELOG.md

@@ -1,3 +1,8 @@
+## 2.17.0-dev
+
+- Adds `companion` entry to `DataClassName` to override the name of the
+  generated companion class.
+
 ## 2.16.0
 
 - When a migration throws, the database will now block subsequent operations

drift/lib/src/dsl/table.dart

@@ -308,7 +308,10 @@ final class TableIndex {
 class DataClassName {
   /// The overridden name to use when generating the data class for a table.
   /// {@macro drift_custom_data_class}
-  final String name;
+  final String? name;
+
+  /// The overridden name to use when generating the companion class for a table.
+  final String? companion;
 
   /// The parent type of the data class generated by drift.
   ///
@@ -345,7 +348,11 @@ class DataClassName {
 
   /// Customize the data class name for a given table.
   /// {@macro drift_custom_data_class}
-  const DataClassName(this.name, {this.extending});
+  const DataClassName(this.name, {this.extending, this.companion});
+
+  /// Customize the data class name for a given table.
+  /// {@macro drift_custom_data_class}
+  const DataClassName.custom({this.name, this.extending, this.companion});
 }
 
 /// An annotation specifying an existing class to be used as a data class.

drift_dev/lib/src/analysis/resolver/dart/helper.dart

@@ -197,12 +197,14 @@ extension TypeUtils on DartType {
 }
 
 class DataClassInformation {
-  final String enforcedName;
+  final String? enforcedName;
+  final String? companionName;
   final CustomParentClass? extending;
   final ExistingRowClass? existingClass;
 
   DataClassInformation(
     this.enforcedName,
+    this.companionName,
     this.extending,
     this.existingClass,
   );
@@ -233,16 +235,15 @@ class DataClassInformation {
       ));
     }
 
-    String name;
+    var name = dataClassName?.getField('name')!.toStringValue();
+    final companionName =
+        dataClassName?.getField('companionName')?.toStringValue();
     CustomParentClass? customParentClass;
     ExistingRowClass? existingClass;
 
     if (dataClassName != null) {
-      name = dataClassName.getField('name')!.toStringValue()!;
       customParentClass =
           parseCustomParentClass(name, dataClassName, element, resolver);
-    } else {
-      name = dataClassNameForClassName(element.name);
     }
 
     if (useRowClass != null) {
@@ -277,7 +278,12 @@ class DataClassInformation {
       }
     }
 
-    return DataClassInformation(name, customParentClass, existingClass);
+    return DataClassInformation(
+      name,
+      companionName,
+      customParentClass,
+      existingClass,
+    );
   }
 }
 

drift_dev/lib/src/analysis/resolver/dart/table.dart

@@ -2,6 +2,7 @@ import 'package:analyzer/dart/ast/ast.dart';
 import 'package:analyzer/dart/ast/syntactic_entity.dart';
 import 'package:analyzer/dart/element/element.dart';
 import 'package:collection/collection.dart';
+import 'package:drift_dev/src/analysis/resolver/shared/data_class.dart';
 import 'package:sqlparser/sqlparser.dart' as sql;
 
 import '../../driver/error.dart';
@@ -54,7 +55,9 @@ class DartTableResolver extends LocalElementResolver<DiscoveredDartTable> {
       DriftDeclaration.dartElement(element),
       columns: columns,
       references: references.toList(),
-      nameOfRowClass: dataClassInfo.enforcedName,
+      nameOfRowClass:
+          dataClassInfo.enforcedName ?? dataClassNameForClassName(element.name),
+      nameOfCompanionClass: dataClassInfo.companionName,
       existingRowClass: dataClassInfo.existingClass,
       customParentClass: dataClassInfo.extending,
       baseDartName: element.name,

drift_dev/lib/src/analysis/resolver/dart/view.dart

@@ -9,6 +9,7 @@ import 'package:recase/recase.dart';
 import '../../results/results.dart';
 import '../intermediate_state.dart';
 import '../resolver.dart';
+import '../shared/data_class.dart';
 import 'helper.dart';
 
 class DartViewResolver extends LocalElementResolver<DiscoveredDartView> {
@@ -26,7 +27,8 @@ class DartViewResolver extends LocalElementResolver<DiscoveredDartView> {
       discovered.ownId,
       DriftDeclaration.dartElement(discovered.dartElement),
       columns: columns,
-      nameOfRowClass: dataClassInfo.enforcedName,
+      nameOfRowClass: dataClassInfo.enforcedName ??
+          dataClassNameForClassName(discovered.dartElement.name),
       existingRowClass: dataClassInfo.existingClass,
       customParentClass: dataClassInfo.extending,
       entityInfoName: '\$${discovered.dartElement.name}View',

drift_dev/lib/src/analysis/resolver/shared/data_class.dart

@@ -31,7 +31,7 @@ String dataClassNameForClassName(String tableName) {
 }
 
 CustomParentClass? parseCustomParentClass(
-  String dartTypeName,
+  String? dartTypeName,
   DartObject dataClassName,
   ClassElement element,
   LocalElementResolver resolver,
@@ -87,7 +87,10 @@ CustomParentClass? parseCustomParentClass(
           code = AnnotatedDartCode([
             DartTopLevelSymbol.topLevelElement(extendingType.element),
             '<',
-            DartTopLevelSymbol(dartTypeName, null),
+            DartTopLevelSymbol(
+              dartTypeName ?? dataClassNameForClassName(element.name),
+              null,
+            ),
             '>',
           ]);
         } else {

drift_dev/lib/src/analysis/results/result_sets.dart

@@ -40,6 +40,9 @@ abstract class DriftElementWithResultSet extends DriftSchemaElement {
   /// The name for the data class associated with this table or view.
   String get nameOfRowClass;
 
+  /// The name for the companion class associated with this table or view.
+  String? get nameOfCompanionClass;
+
   /// All [columns] of this table, indexed by their name in SQL.
   late final Map<String, DriftColumn> columnBySqlName = CaseInsensitiveMap.of({
     for (final column in columns) column.nameInSql: column,

drift_dev/lib/src/analysis/results/table.dart

@@ -32,6 +32,9 @@ class DriftTable extends DriftElementWithResultSet {
   @override
   final String nameOfRowClass;
 
+  @override
+  final String? nameOfCompanionClass;
+
   final bool withoutRowId;
 
   /// Information about the virtual table creating statement backing this table,
@@ -69,6 +72,7 @@ class DriftTable extends DriftElementWithResultSet {
     required this.columns,
     required this.baseDartName,
     required this.nameOfRowClass,
+    this.nameOfCompanionClass,
     this.references = const [],
     this.existingRowClass,
     this.customParentClass,

drift_dev/lib/src/analysis/results/view.dart

@@ -25,6 +25,9 @@ class DriftView extends DriftElementWithResultSet {
   @override
   final String nameOfRowClass;
 
+  @override
+  final String? nameOfCompanionClass;
+
   @override
   List<DriftElement> references;
 
@@ -38,6 +41,7 @@ class DriftView extends DriftElementWithResultSet {
     required this.existingRowClass,
     required this.nameOfRowClass,
     required this.references,
+    this.nameOfCompanionClass,
   });
 
   @override

drift_dev/lib/src/analysis/serializer.dart

@@ -57,6 +57,7 @@ class ElementSerializer {
         'fixed_entity_info_name': element.fixedEntityInfoName,
         'base_dart_name': element.baseDartName,
         'row_class_name': element.nameOfRowClass,
+        'companion_class_name': element.nameOfCompanionClass,
         'without_rowid': element.withoutRowId,
         'strict': element.strict,
         if (element.isVirtual)
@@ -146,6 +147,7 @@ class ElementSerializer {
         'custom_parent_class':
             _serializeCustomParentClass(element.customParentClass),
         'name_of_row_class': element.nameOfRowClass,
+        'name_of_companion_class': element.nameOfCompanionClass,
         'source': serializedSource,
       };
     } else if (element is BaseDriftAccessor) {
@@ -536,6 +538,7 @@ class ElementDeserializer {
           fixedEntityInfoName: json['fixed_entity_info_name'] as String?,
           baseDartName: json['base_dart_name'] as String,
           nameOfRowClass: json['row_class_name'] as String,
+          nameOfCompanionClass: json['companion_class_name'] as String?,
           withoutRowId: json['without_rowid'] as bool,
           strict: json['strict'] as bool,
           virtualTableData: virtualTableData,
@@ -675,6 +678,7 @@ class ElementDeserializer {
           customParentClass:
               _readCustomParentClass(json['custom_parent_class'] as Map?),
           nameOfRowClass: json['name_of_row_class'] as String,
+          nameOfCompanionClass: json['name_of_companion_class'] as String,
           existingRowClass: json['existing_data_class'] != null
               ? await _readExistingRowClass(
                   id.libraryUri, json['existing_data_class'] as Map)

drift_dev/lib/src/writer/writer.dart

@@ -72,9 +72,10 @@ abstract class _NodeOrWriter {
   }
 
   AnnotatedDartCode companionType(DriftTable table) {
-    final baseName = writer.options.useDataClassNameForCompanions
-        ? table.nameOfRowClass
-        : table.baseDartName;
+    final baseName = table.nameOfCompanionClass ??
+        (writer.options.useDataClassNameForCompanions
+            ? table.nameOfRowClass
+            : table.baseDartName);
 
     return generatedElement(table, '${baseName}Companion');
   }

drift_dev/test/writer/data_class_writer_test.dart

@@ -319,6 +319,68 @@ class Database extends _$Database {}
     },
     tags: 'analyzer',
   );
+
+  test(
+    'It should use the provided names for the data classes and the companion class',
+    () async {
+      final writer = await emulateDriftBuild(
+        inputs: const {
+          'a|lib/main.dart': r'''
+import 'package:drift/drift.dart';
+
+part 'main.drift.dart';
+
+@DataClassName('FirstDataClass', companion: 'FirstCompanionClass')
+class FirstTable extends Table {
+  TextColumn get foo => text()();
+  IntColumn get bar => integer()();
+}
+
+@DataClassName.custom(name: 'SecondDataClass', companion: 'SecondCompanionClass')
+class SecondTable extends Table {
+  TextColumn get foo => text()();
+  IntColumn get bar => integer()();
+}
+
+@DataClassName.custom(companion: 'ThirdCompanionClass')
+class ThirdTable extends Table {
+  TextColumn get foo => text()();
+  IntColumn get bar => integer()();
+}
+
+@DriftDatabase(
+  tables: [FirstTable, SecondTable, ThirdTable],
+)
+class Database extends _$Database {}
+'''
+        },
+      );
+
+      checkOutputs({
+        'a|lib/main.drift.dart': decodedMatches(allOf([
+          contains(
+            'class FirstDataClass extends DataClass implements Insertable<FirstDataClass> {',
+          ),
+          contains(
+            'class FirstTableCompanion extends UpdateCompanion<FirstDataClass> {',
+          ),
+          contains(
+            'class SecondDataClass extends DataClass implements Insertable<SecondDataClass> {',
+          ),
+          contains(
+            'class SecondTableCompanion extends UpdateCompanion<SecondDataClass> {',
+          ),
+          contains(
+            'class ThirdTableData extends DataClass implements Insertable<ThirdTableData> {',
+          ),
+          contains(
+            'class ThirdTableCompanion extends UpdateCompanion<ThirdTableData> {',
+          ),
+        ])),
+      }, writer.dartOutputs, writer.writer);
+    },
+    tags: 'analyzer',
+  );
 }
 
 class _GeneratesConstDataClasses extends Matcher {