Fix enums from absolute imports

Author: guaycuru

pydantic2ts/cli/script.py

@@ -118,21 +118,16 @@ def extract_pydantic_models(module: ModuleType) -> List[Type[BaseModel]]:
     return models
 
 
-def extract_enum_models(module: ModuleType) -> List[Type[Enum]]:
+def extract_enum_models(models: List[Type[BaseModel]]) -> List[Type[Enum]]:
     """
-    Given a module, return a list of the Enum classes contained within it.
+    Given a list of pydantic models, return a list of the Enum classes used as fields within those models.
     """
     enums = []
-    module_name = module.__name__
-
-    for _, enum in inspect.getmembers(module, is_enum):
-        if enum.__module__ != "enum":
-            enums.append(enum)
 
-    for _, submodule in inspect.getmembers(
-        module, lambda obj: is_submodule(obj, module_name)
-    ):
-        enums.extend(extract_enum_models(submodule))
+    for model in models:
+        for field in model.__fields__.values():
+            if isinstance(field.type_, type) and issubclass(field.type_, Enum):
+                enums.append(field.type_)
 
     return enums
 
@@ -315,8 +310,7 @@ def generate_typescript_defs(
 
     logger.info("Finding pydantic models...")
 
-    import_result = import_module(module)
-    models = extract_pydantic_models(import_result)
+    models = extract_pydantic_models(import_module(module))
 
     if exclude:
         models = [m for m in models if m.__name__ not in exclude]
@@ -326,7 +320,7 @@ def generate_typescript_defs(
     if V2:
         schema = generate_json_schema_v2(models)
     else:
-        enums = extract_enum_models(import_result)
+        enums = extract_enum_models(models)
         schema = generate_json_schema_v1(models, enums)
 
     schema_dir = mkdtemp()

tests/expected_results/enums/v1/input.py

@@ -1,10 +1,16 @@
+import os
+import sys
 from enum import Enum
 from typing import List, Optional
 
 from pydantic import BaseModel
 
+# Make absolute imports work
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__))))
+
 from .schemas.schema_one import ModelOne  # noqa: F401
 from .schemas.schema_two import ModelTwo  # noqa: F401
+from schemas.sub_model import SubModel  # this tests absolute imports
 
 
 class Cat(BaseModel):
@@ -29,3 +35,7 @@ class AnimalShelter(BaseModel):
 class Standalone(Enum):
     something = "something"
     anything = "anything"
+
+
+class ImportedSubModule(BaseModel):
+    sub: SubModel

tests/expected_results/enums/v1/output.ts

@@ -21,13 +21,23 @@ export interface Dog {
   name: string;
   age: number;
 }
+export interface ImportedSubModule {
+  sub: SubModel;
+}
+export interface SubModel {
+  bar: Bar;
+}
 export interface ModelOne {
   foo: Foo;
 }
 export interface ModelTwo {
   foo: Foo1;
 }
 
+export const enum Bar {
+  ONE = "one",
+  TWO = "two"
+}
 export const enum Foo {
   ONE_A = "one_a",
   ONE_B = "one_b"

tests/expected_results/enums/v1/schemas/sub_model.py

@@ -0,0 +1,12 @@
+import enum
+
+from pydantic import BaseModel
+
+
+class Bar(enum.Enum):
+    ONE = "one"
+    TWO = "two"
+
+
+class SubModel(BaseModel):
+    bar: Bar

tests/expected_results/enums/v2/input.py

@@ -1,10 +1,16 @@
+import os
+import sys
 from enum import Enum
 from typing import List, Optional
 
 from pydantic import BaseModel
 
+# Make absolute imports work
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__))))
+
 from .schemas.schema_one import ModelOne  # noqa: F401
 from .schemas.schema_two import ModelTwo  # noqa: F401
+from schemas.sub_model import SubModel  # this tests absolute imports
 
 
 class Cat(BaseModel):
@@ -29,3 +35,7 @@ class AnimalShelter(BaseModel):
 class Standalone(Enum):
     something = "something"
     anything = "anything"
+
+
+class ImportedSubModule(BaseModel):
+    sub: SubModel

tests/expected_results/enums/v2/output.ts

@@ -21,13 +21,23 @@ export interface Dog {
   name: string;
   age: number;
 }
+export interface ImportedSubModule {
+  sub: SubModel;
+}
+export interface SubModel {
+  bar: Bar;
+}
 export interface ModelOne {
   foo: Foo;
 }
 export interface ModelTwo {
   foo: Foo1;
 }
 
+export const enum Bar {
+  ONE = "one",
+  TWO = "two"
+}
 export const enum Foo {
   ONE_A = "one_a",
   ONE_B = "one_b"

tests/expected_results/enums/v2/schemas/sub_model.py

@@ -0,0 +1,12 @@
+import enum
+
+from pydantic import BaseModel
+
+
+class Bar(enum.Enum):
+    ONE = "one"
+    TWO = "two"
+
+
+class SubModel(BaseModel):
+    bar: Bar