From c1ee1f7274d00b39cb01d38dd035b227c1a87916 Mon Sep 17 00:00:00 2001
From: Alexander Wilde <alexander.wilde87@gmail.com>
Date: Thu, 1 Dec 2022 02:04:16 +0000
Subject: [PATCH] Add JsonTypeConverter support to enums

---
 drift/lib/src/runtime/types/converters.dart                | 6 ++++--
 drift_dev/lib/src/analysis/resolver/dart/column.dart       | 5 ++++-
 drift_dev/lib/src/analysis/resolver/drift/table.dart       | 2 ++
 drift_dev/lib/src/analysis/resolver/shared/dart_types.dart | 6 +++++-
 drift_dev/lib/src/writer/tables/data_class_writer.dart     | 5 ++++-
 5 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/drift/lib/src/runtime/types/converters.dart b/drift/lib/src/runtime/types/converters.dart
index febc0499..5ce18b35 100644
--- a/drift/lib/src/runtime/types/converters.dart
+++ b/drift/lib/src/runtime/types/converters.dart
@@ -83,7 +83,8 @@ mixin JsonTypeConverter<D, S> implements JsonTypeConverter2<D, S, S> {
 
 /// Implementation for an enum to int converter that uses the index of the enum
 /// as the value stored in the database.
-class EnumIndexConverter<T extends Enum> extends TypeConverter<T, int> {
+class EnumIndexConverter<T extends Enum> extends TypeConverter<T, int>
+    with JsonTypeConverter<T, int> {
   /// All values of the enum.
   final List<T> values;
 
@@ -103,7 +104,8 @@ class EnumIndexConverter<T extends Enum> extends TypeConverter<T, int> {
 
 /// Implementation for an enum to string converter that uses the name of the
 /// enum as the value stored in the database.
-class EnumNameConverter<T extends Enum> extends TypeConverter<T, String> {
+class EnumNameConverter<T extends Enum> extends TypeConverter<T, String>
+    with JsonTypeConverter<T, String> {
   /// All values of the enum.
   final List<T> values;
 
diff --git a/drift_dev/lib/src/analysis/resolver/dart/column.dart b/drift_dev/lib/src/analysis/resolver/dart/column.dart
index 75f38124..9ddc9c99 100644
--- a/drift_dev/lib/src/analysis/resolver/dart/column.dart
+++ b/drift_dev/lib/src/analysis/resolver/dart/column.dart
@@ -337,6 +337,7 @@ class ColumnParser {
 
     final sqlName = foundExplicitName ?? ReCase(getter.name.lexeme).snakeCase;
     final sqlType = _startMethodToColumnType(foundStartMethod);
+    final helper = await _resolver.resolver.driver.loadKnownTypes();
 
     AppliedTypeConverter? converter;
     if (mappedAs != null) {
@@ -347,7 +348,7 @@ class ColumnParser {
         nullable,
         (message) => _resolver.reportError(
             DriftAnalysisError.inDartAst(element, mappedAs!, message)),
-        await _resolver.resolver.driver.loadKnownTypes(),
+        helper,
       );
     }
 
@@ -366,6 +367,7 @@ class ColumnParser {
             remainingExpr.typeArguments ?? remainingExpr.methodName, msg)),
         enumType,
         EnumType.intEnum,
+        helper,
       );
     } else if (foundStartMethod == _startTextEnum) {
       if (converter != null) {
@@ -382,6 +384,7 @@ class ColumnParser {
             remainingExpr.typeArguments ?? remainingExpr.methodName, msg)),
         enumType,
         EnumType.textEnum,
+        helper,
       );
     }
 
diff --git a/drift_dev/lib/src/analysis/resolver/drift/table.dart b/drift_dev/lib/src/analysis/resolver/drift/table.dart
index e78cc001..61a42003 100644
--- a/drift_dev/lib/src/analysis/resolver/drift/table.dart
+++ b/drift_dev/lib/src/analysis/resolver/drift/table.dart
@@ -28,6 +28,7 @@ class DriftTableResolver extends LocalElementResolver<DiscoveredDriftTable> {
     Table table;
     final references = <DriftElement>{};
     final stmt = discovered.sqlNode;
+    final helper = await resolver.driver.loadKnownTypes();
 
     try {
       final reader = SchemaFromCreateTable(
@@ -74,6 +75,7 @@ class DriftTableResolver extends LocalElementResolver<DiscoveredDriftTable> {
                 DriftAnalysisError.inDriftFile(column.definition ?? stmt, msg)),
             dartClass.classElement.thisType,
             type == DriftSqlType.int ? EnumType.intEnum : EnumType.textEnum,
+            helper,
           );
         }
       }
diff --git a/drift_dev/lib/src/analysis/resolver/shared/dart_types.dart b/drift_dev/lib/src/analysis/resolver/shared/dart_types.dart
index 0f251bf8..8703eee4 100644
--- a/drift_dev/lib/src/analysis/resolver/shared/dart_types.dart
+++ b/drift_dev/lib/src/analysis/resolver/shared/dart_types.dart
@@ -230,7 +230,9 @@ AppliedTypeConverter readEnumConverter(
   void Function(String) reportError,
   DartType dartEnumType,
   EnumType columnEnumType,
+  KnownDriftTypes helper,
 ) {
+  final typeProvider = helper.helperLibrary.typeProvider;
   if (dartEnumType is! InterfaceType) {
     reportError('Not a class: `$dartEnumType`');
   }
@@ -261,7 +263,9 @@ AppliedTypeConverter readEnumConverter(
   return AppliedTypeConverter(
     expression: expression,
     dartType: dartEnumType,
-    jsonType: null,
+    jsonType: columnEnumType == EnumType.intEnum
+        ? typeProvider.intType
+        : typeProvider.stringType,
     sqlType: columnEnumType == EnumType.intEnum
         ? DriftSqlType.int
         : DriftSqlType.string,
diff --git a/drift_dev/lib/src/writer/tables/data_class_writer.dart b/drift_dev/lib/src/writer/tables/data_class_writer.dart
index e63ecd4a..e9707e83 100644
--- a/drift_dev/lib/src/writer/tables/data_class_writer.dart
+++ b/drift_dev/lib/src/writer/tables/data_class_writer.dart
@@ -133,8 +133,11 @@ class DataClassWriter {
 
       final typeConverter = column.typeConverter;
       if (typeConverter != null && typeConverter.alsoAppliesToJsonConversion) {
-        final type =
+        var type =
             _emitter.dartCode(AnnotatedDartCode.type(typeConverter.jsonType!));
+        if (column.nullable) {
+          type = '$type?';
+        }
 
         final fromConverter = "serializer.fromJson<$type>(json['$jsonKey'])";
         final converterField = _converter(column);