add typing_extensions.Literals

devdoomari3 · devdoomari3 · commit 280486cfbfb4 · 2019-05-14T23:52:29.000+09:00
diff --git a/py_codegen/test_fixtures/func_with_literals.py b/py_codegen/test_fixtures/func_with_literals.py
@@ -0,0 +1,11 @@
+from typing_extensions import Literal
+
+
+def func_with_literals(
+        input1: Literal[
+            'a',
+            1,
+            Literal[2,3],
+        ],
+) -> Literal[True, 0.5, Literal[3]]:
+    return input1
diff --git a/py_codegen/type_extractor/__tests__/test_func_with_literals.py b/py_codegen/type_extractor/__tests__/test_func_with_literals.py
@@ -0,0 +1,41 @@
+from py_codegen.type_extractor.__tests__.utils import cleanup, traverse
+from py_codegen.type_extractor.nodes.FunctionFound import FunctionFound
+from py_codegen.type_extractor.nodes.LiteralFound import LiteralFound
+from py_codegen.type_extractor.nodes.TypeOR import TypeOR
+from py_codegen.type_extractor.type_extractor import TypeExtractor
+from py_codegen.test_fixtures.func_with_literals import func_with_literals
+
+
+def test_func_with_list():
+    type_collector = TypeExtractor()
+
+    type_collector.add(None)(func_with_literals)
+
+    func_found_cleaned = cleanup(
+        type_collector.collected_types[func_with_literals.__qualname__],
+    )
+    assert func_found_cleaned == traverse(
+        FunctionFound(
+            name=func_with_literals.__qualname__,
+            params={
+                'input1': TypeOR(
+                    a=LiteralFound('a'),
+                    b=TypeOR(
+                        a=LiteralFound(1),
+                        b=TypeOR(
+                            a=LiteralFound(2),
+                            b=LiteralFound(3),
+                        ),
+                    ),
+                ),
+            },
+            return_type=TypeOR(
+                a=LiteralFound(True),
+                b=TypeOR(
+                    a=LiteralFound(0.5),
+                    b=LiteralFound(3),
+                ),
+            )
+        ),
+        cleanup,
+    )
diff --git a/py_codegen/type_extractor/middlewares/literal_found.py b/py_codegen/type_extractor/middlewares/literal_found.py
@@ -0,0 +1,25 @@
+from typing import Tuple
+from typing_extensions import Literal
+
+from py_codegen.type_extractor.__base__ import BaseTypeExtractor
+from py_codegen.type_extractor.nodes.LiteralFound import LiteralFound
+from py_codegen.type_extractor.nodes.TypeOR import TypeOR
+
+
+def literal_found_middleware(typ, type_extractor: BaseTypeExtractor):
+    typ_origin = typ.__origin__
+    if typ_origin is not Literal:
+        return
+    return __process_literal_args(typ.__args__)
+
+
+def __process_literal_args(args: Tuple):
+    current = LiteralFound(args[0])
+    try:
+        if args[0].__origin__ is Literal:
+            current = __process_literal_args(args[0].__args__)
+    except:
+        pass
+    if len(args) == 1:
+        return current
+    return TypeOR(current, __process_literal_args(args[1:]))
diff --git a/py_codegen/type_extractor/nodes/LiteralFound.py b/py_codegen/type_extractor/nodes/LiteralFound.py
@@ -0,0 +1,7 @@
+from typing import NamedTuple, Any
+
+from py_codegen.type_extractor.nodes.BaseNodeType import BaseNodeType
+
+
+class LiteralFound(NamedTuple, BaseNodeType):  # type: ignore
+    value: Any
diff --git a/py_codegen/type_extractor/type_extractor.py b/py_codegen/type_extractor/type_extractor.py
@@ -16,6 +16,7 @@
 from py_codegen.type_extractor.middlewares.dict_found import dict_found_middleware
 from py_codegen.type_extractor.middlewares.function_found import func_found_middleware
 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.tuple_found import tuple_found_middleware
 from py_codegen.type_extractor.middlewares.type_or import typeor_middleware
@@ -44,6 +45,7 @@ class TypeExtractor(BaseTypeExtractor):
         list_found_middleware,
         typeor_middleware,
         typeddict_found_middleware,
+        literal_found_middleware,
         dict_found_middleware,
         tuple_found_middleware,
         class_found_middleware,
