Make extension references explicit

drift_dev/CHANGELOG.md

@@ -1,11 +1,12 @@
-## 2.18.1-dev
+## 2.19.0-dev
 
 - Fix generated `CREATE VIEW` statements containing existing row class syntax
   only supposed to be used during static analysis.
 - Fix Dart views referencing the same column from different table aliases using
   columns with the same name.
 - Fix `drift_dev schema steps` generating invalid code when no migrations have
   been defined yet.
+- Fix generated imports for extension member references in modular mode.
 
 ## 2.18.0
 

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

@@ -16,6 +16,10 @@ import 'types.dart';
 
 part '../../generated/analysis/results/dart.g.dart';
 
+/// A syntactic representation of Dart code where top-level symbols (that would)
+/// have to be imported are explicitly represented with their defining URL.
+///
+/// This allows generating code with independent import management.
 class AnnotatedDartCode {
   static final Uri dartAsync = Uri.parse('dart:async');
   static final Uri dartCore = Uri.parse('dart:core');
@@ -509,6 +513,37 @@ class _AddFromAst extends GeneralizingAstVisitor<void> {
     node.argumentList.accept(this);
   }
 
+  @override
+  void visitMethodInvocation(MethodInvocation node) {
+    // Rewrite extension invocations (e.g. `myList.indexed`) to explicitly
+    // mention the extension in use (e.g `IterableExtensions(myList).indexed`).
+    // This is because extensions aren't visible as soon as the library defining
+    // them is imported under an import alias. Explicitly mentioning the
+    // extension fixes this problem.
+    if (node.target == null || node.realTarget != node.target) {
+      // Unfortunately there's no easy way to apply this to cascade expressions
+      return super.visitMethodInvocation(node);
+    }
+
+    final element = node.methodName.staticElement;
+    final enclosing = element?.enclosingElement;
+    if (enclosing is! ExtensionElement || enclosing.name == null) {
+      return super.visitMethodInvocation(node);
+    }
+
+    _builder
+      ..addTopLevel(
+          DartTopLevelSymbol.topLevelElement(enclosing, enclosing.name!))
+      ..addText('(');
+    node.target?.accept(this);
+    _builder
+      ..addText(node.isNullAware ? ')?.' : ').')
+      ..addText(node.methodName.name);
+
+    node.typeArguments?.accept(this);
+    node.argumentList.accept(this);
+  }
+
   @override
   void visitNamedType(NamedType node) {
     _addTopLevelReference(node.element, node.name2);

drift_dev/test/analysis/results/dart_test.dart

@@ -8,7 +8,13 @@ import '../test_utils.dart';
 void main() {
   late TestBackend tester;
 
-  setUpAll(() async => tester = await TestBackend.init({}));
+  setUpAll(() async => tester = await TestBackend.init({
+        'a|lib/definitions.dart': '''
+extension MyStringUtils on String {
+  String reverse() => throw 'todo';
+}
+''',
+      }));
   tearDownAll(() => tester.dispose());
 
   group('from AST', () {
@@ -17,8 +23,8 @@ void main() {
     Future<void> checkTransformation(String sourceExpression,
         String expectedResult, Map<String, String> expectedImports) async {
       final testUri = Uri.parse('package:a/test_${testCount++}.dart');
-      final expression =
-          await tester.resolveExpression(testUri, sourceExpression, const []);
+      final expression = await tester.resolveExpression(
+          testUri, sourceExpression, const ['package:a/definitions.dart']);
       final annotated = AnnotatedDartCode.ast(expression);
 
       final imports = TestImportManager();
@@ -53,5 +59,19 @@ void main() {
           'i0.IterableExtensions<i1.String>([]).firstOrNull',
           {'i0': 'dart:collection', 'i1': 'dart:core'});
     });
+
+    test('extension method invocations', () async {
+      await checkTransformation(
+        "'hello world'.reverse()",
+        "i0.MyStringUtils('hello world').reverse()",
+        {'i0': 'package:a/definitions.dart'},
+      );
+
+      await checkTransformation(
+        "'hello world'?.reverse<void>(1, 2, 3)",
+        "i0.MyStringUtils('hello world')?.reverse<void>(1,2,3)",
+        {'i0': 'package:a/definitions.dart'},
+      );
+    });
   });
 }