add newtype support

devdoomari3 · devdoomari3 · commit 6f555f3d193c · 2020-05-10T17:44:05.000+09:00
diff --git a/py_codegen/test_fixtures/new_type.py b/py_codegen/test_fixtures/new_type.py
@@ -0,0 +1,17 @@
+from typing import (
+    NewType,
+)
+
+
+UserId = NewType('UserId', int)
+
+
+def name_by_id(user_id: UserId) -> str:
+    return str(user_id)
+
+
+class SomeClass:
+    a: int
+
+
+SomeClassInNewType = NewType('SomeClassInNewType', SomeClass)
diff --git a/py_codegen/type_extractor/__tests__/test_newtype.py b/py_codegen/type_extractor/__tests__/test_newtype.py
@@ -0,0 +1,46 @@
+from py_codegen.test_fixtures.new_type import (
+    name_by_id,
+    SomeClassInNewType,
+)
+from py_codegen.type_extractor.__tests__.utils import traverse, cleanup
+from py_codegen.type_extractor.nodes.ClassFound import ClassFound
+from py_codegen.type_extractor.nodes.FunctionFound import FunctionFound
+from py_codegen.type_extractor.nodes.NewType import NewTypeFound
+from py_codegen.type_extractor.nodes.TypeVarFound import TypeVarFound
+from py_codegen.type_extractor.type_extractor import TypeExtractor
+
+
+def test_newtype():
+    type_extractor = TypeExtractor()
+    type_extractor.add()(name_by_id)
+    type_extractor.add()(SomeClassInNewType)
+    cleaned_up = {
+        key: traverse(value, cleanup)
+        for (key, value) in type_extractor.collected_types.items()
+    }
+    user_id_newtype = NewTypeFound(
+        name='UserId',
+        actual=int,
+    )
+    some_class = ClassFound(
+        name='SomeClass',
+        fields={
+            'a': int,
+        }
+    )
+    assert cleaned_up == {
+        'UserId': user_id_newtype,
+        'name_by_id': FunctionFound(
+            name='name_by_id',
+            params={
+                'user_id': user_id_newtype,
+            },
+            return_type=str,
+        ),
+        'SomeClass': some_class,
+        'SomeClassInNewType': NewTypeFound(
+            name='SomeClassInNewType',
+            actual=some_class,
+        ),
+    }
+    print(1)
diff --git a/py_codegen/type_extractor/__tests__/utils.py b/py_codegen/type_extractor/__tests__/utils.py
@@ -4,6 +4,7 @@
 
 from py_codegen.type_extractor.nodes.BaseNodeType import NodeType
 from py_codegen.type_extractor.nodes.FixedGenericFound import FixedGenericFound
+from py_codegen.type_extractor.nodes.NewType import NewTypeFound
 from py_codegen.type_extractor.nodes.TypeVarFound import TypeVarFound
 from py_codegen.type_extractor.nodes.TypedDictFound import TypedDictFound
 from py_codegen.type_extractor.nodes.ClassFound import ClassFound
@@ -67,6 +68,12 @@ def traverse(node: NodeType, func: traverse_func_type):
             node.original, func
         )
         return func(typevar_node)
+    if isinstance(node, NewTypeFound):
+        newtype_found = copy(node)
+        newtype_found.actual = traverse(
+            node.actual, func
+        )
+        return func(newtype_found)
     return node
 
 
diff --git a/py_codegen/type_extractor/middlewares/newtype_found.py b/py_codegen/type_extractor/middlewares/newtype_found.py
@@ -0,0 +1,28 @@
+import inspect
+
+from typing import Set
+from py_codegen.type_extractor.__base__ import BaseTypeExtractor
+
+from py_codegen.type_extractor.nodes.BaseNodeType import BaseOption
+from py_codegen.type_extractor.nodes.NewType import NewTypeFound
+
+
+def newtype_found_middleware(
+        typ,
+        type_extractor: BaseTypeExtractor,
+        options: Set[BaseOption],
+):
+    if not inspect.isfunction(typ):
+        return None
+    if typ.__module__ is not 'typing':
+        return None
+    if typ.__code__.co_name is not 'new_type':
+        return None
+
+    newtype_found = NewTypeFound(
+        name=typ.__name__,
+        actual=type_extractor.rawtype_to_node(typ.__supertype__),
+    )
+
+    type_extractor.collected_types[newtype_found.name] = newtype_found
+    return newtype_found
diff --git a/py_codegen/type_extractor/nodes/NewType.py b/py_codegen/type_extractor/nodes/NewType.py
@@ -0,0 +1,10 @@
+from dataclasses import dataclass
+from typing import List
+
+from py_codegen.type_extractor.nodes.BaseNodeType import BaseNodeType, NodeType
+
+
+@dataclass
+class NewTypeFound(BaseNodeType):
+    name: str
+    actual: NodeType
diff --git a/py_codegen/type_extractor/type_extractor.py b/py_codegen/type_extractor/type_extractor.py
@@ -19,6 +19,7 @@
 from py_codegen.type_extractor.middlewares.list_found import list_found_middleware
 from py_codegen.type_extractor.middlewares.literal_found import literal_found_middleware
 from py_codegen.type_extractor.middlewares.mapping_found import mapping_found_middleware
+from py_codegen.type_extractor.middlewares.newtype_found import newtype_found_middleware
 from py_codegen.type_extractor.middlewares.tuple_found import tuple_found_middleware
 from py_codegen.type_extractor.middlewares.type_or import typeor_middleware
 from py_codegen.type_extractor.nodes.BaseNodeType import NodeType, BaseNodeType, BaseOption
@@ -54,6 +55,7 @@ class TypeExtractor(BaseTypeExtractor):
         typevar_found_middleware,
         fixed_generic_found_middleware,
         class_found_middleware,
+        newtype_found_middleware,
         func_found_middleware,
         none_node_middleware,
         builtin_middleware,
