Use mypy's make_optional_type instead of shipping our own (#2334)

typeddjango · Aug 10, 2024 · 3e6f1a6 · 3e6f1a6
1 parent d4d7be7
commit 3e6f1a6
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 7 deletions.
diff --git a/mypy_django_plugin/django/context.py b/mypy_django_plugin/django/context.py
@@ -31,6 +31,7 @@
 from mypy.checker import TypeChecker
 from mypy.nodes import TypeInfo
 from mypy.plugin import MethodContext
+from mypy.typeanal import make_optional_type
 from mypy.types import AnyType, Instance, ProperType, TypeOfAny, UnionType, get_proper_type
 from mypy.types import Type as MypyType
 
@@ -182,7 +183,7 @@ def get_field_lookup_exact_type(
             primary_key_type = self.get_field_get_type(api, rel_model_info, primary_key_field, method="init")
 
             model_and_primary_key_type = UnionType.make_union([Instance(rel_model_info, []), primary_key_type])
-            return helpers.make_optional(model_and_primary_key_type)
+            return make_optional_type(model_and_primary_key_type)
 
         field_info = helpers.lookup_class_typeinfo(api, field.__class__)
         if field_info is None:

diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py
@@ -36,6 +36,7 @@
     MethodContext,
 )
 from mypy.semanal import SemanticAnalyzer
+from mypy.typeanal import make_optional_type
 from mypy.types import (
     AnyType,
     Instance,
@@ -232,10 +233,6 @@ def get_call_argument_type_by_name(ctx: Union[FunctionContext, MethodContext], n
     return arg_types[0]
 
 
-def make_optional(typ: MypyType) -> MypyType:
-    return UnionType.make_union([typ, NoneTyp()])
-
-
 def is_optional(typ: MypyType) -> bool:
     typ = get_proper_type(typ)
     return isinstance(typ, UnionType) and any(isinstance(get_proper_type(item), NoneTyp) for item in typ.items)
@@ -277,7 +274,7 @@ def get_private_descriptor_type(type_info: TypeInfo, private_field_name: str, is
             return AnyType(TypeOfAny.explicit)
 
         if is_nullable:
-            descriptor_type = make_optional(descriptor_type)
+            descriptor_type = make_optional_type(descriptor_type)
         return descriptor_type
     return AnyType(TypeOfAny.explicit)
 
@@ -289,7 +286,7 @@ def get_field_lookup_exact_type(api: TypeChecker, field: "Field[Any, Any]") -> M
         rel_model_info = lookup_class_typeinfo(api, lookup_type_class)
         if rel_model_info is None:
             return AnyType(TypeOfAny.from_error)
-        return make_optional(Instance(rel_model_info, []))
+        return make_optional_type(Instance(rel_model_info, []))
 
     field_info = lookup_class_typeinfo(api, field.__class__)
     if field_info is None: