Add support for making external types draftable

packages/draft/README.md

@@ -143,6 +143,24 @@ final c = C([1, 2, 3]).produce((draft) {
 
 Draft is unopinionated and does not provide any sort of equality checking out of the box. If you want equality checking, consider using [equatable](https://pub.dev/packages/equatable)
 
+### External classes
+
+You can make classes from external packages draftable by annotating an extension on them with `@draft`:
+
+```dart
+import 'package:foo_class/foo_class.dart' show Foo;
+
+@draft
+extension on Foo {}
+```
+Once annotated, the external class behaves like any other draftable type and supports all Draft features:
+
+```dart
+Foo(value: 1).produce((draft) {
+  draft.value += 1;
+});
+```
+
 ## Contributing
 
 If you like the package and want to contribute, feel free to [open and issue or create a PR](https://github.com/josiahsrc/draft/tree/main). I'm always open to suggestions and improvements.

packages/draft_builder/lib/src/generator.dart

@@ -476,10 +476,17 @@ class DraftGenerator extends GeneratorForAnnotation<Draft> {
     ConstantReader annotation,
     BuildStep buildStep,
   ) {
+    final classElement = switch (element) {
+      // Direct class annotation.
+      ClassElement element => element,
+      // Extension on a class.
+      ExtensionElement(extendedType: DartType(:ClassElement element)) => element,
+      _ => null,
+    };
+
     // Process only classes.
-    if (element is! ClassElement) return '';
+    if (classElement == null) return '';
 
-    final classElement = element;
     final className = classElement.name;
     final draftClassName = _draftTypeName(className);
 

packages/draft_builder/test/extension_test.dart

@@ -0,0 +1,37 @@
+import 'package:test/test.dart';
+
+import 'common.dart';
+import 'integration/extension.dart';
+
+void main() {
+  test('compiles', () async {
+    await expectLater(
+      compile(r'''
+import 'extension.dart';
+
+void main() {
+  BasicWithExtension(value: 1);
+  BasicWithExtension(value: 1).draft().save();
+  BasicWithExtension(value: 1).produce((draft) {
+    draft.value = 2;
+  });
+}
+'''),
+      completes,
+    );
+  });
+
+  test('works correctly', () async {
+    expect(BasicWithExtension(value: 1).draft().value, 1);
+    expect(BasicWithExtension(value: 1).draft().save().value, 1);
+    expect(
+      BasicWithExtension(value: 1).produce((draft) {
+        draft.value = 2;
+      }).value,
+      2,
+    );
+
+    final draft = BasicWithExtension(value: 1).draft()..value = 10;
+    expect(draft.save().value, 10);
+  });
+}

packages/draft_builder/test/integration/extension.dart

@@ -0,0 +1,12 @@
+import 'package:draft/draft.dart';
+
+part 'extension.draft.dart';
+
+@draft
+extension on BasicWithExtension {}
+
+class BasicWithExtension {
+  final int value;
+
+  BasicWithExtension({required this.value});
+}