[cdd/{emit/utils/sqlalchemy_utils.py,parse/utils/json_schema_utils.py…

…,tests/test_emit/test_emit_sqlalchemy_utils.py}] Use correct `Union` syntax https://docs.python.org/3/library/typing.html#typing.Union
offscale · Jan 9, 2023 · 13ad65f · 13ad65f
1 parent cd15d9a
commit 13ad65f
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 11 deletions.
diff --git a/cdd/emit/utils/sqlalchemy_utils.py b/cdd/emit/utils/sqlalchemy_utils.py
@@ -7,16 +7,18 @@
     AST,
     Assign,
     Attribute,
-    BinOp,
     Call,
     ClassDef,
     Expr,
     FunctionDef,
     ImportFrom,
+    Index,
     Load,
     Module,
     Name,
     Return,
+    Subscript,
+    Tuple,
     alias,
     arguments,
 )
@@ -198,15 +200,13 @@ def update_args_infer_typ_sqlalchemy(_param, args, name, nullable, x_typ_sql):
     elif _param.get("typ").startswith("Union["):
         # Hack to remove the union type. Enum parse seems to be incorrect?
         union_typ = ast.parse(_param["typ"]).body[0]
-        assert isinstance(
-            union_typ.value.slice, BinOp
-        ), "Expected `BinOp` got `{type_name}`".format(
-            type_name=type(union_typ.value.slice).__name__
-        )
-        left, right = map(
-            rpartial(str.rstrip, "\n"),
-            map(to_code, (union_typ.value.slice.left, union_typ.value.slice.right)),
+        assert (
+            isinstance(union_typ.value, Subscript)
+            and isinstance(union_typ.value.slice, Index)
+            and isinstance(union_typ.value.slice.value, Tuple)
         )
+        assert len(union_typ.value.slice.value.elts) == 2
+        left, right = map(attrgetter("id"), union_typ.value.slice.value.elts)
         args.append(
             Name(
                 typ2column_type.get(right, right)

diff --git a/cdd/parse/utils/json_schema_utils.py b/cdd/parse/utils/json_schema_utils.py
@@ -76,7 +76,7 @@ def transform_ref_fk_set(ref, foreign_key):
         _param["typ"] = (
             _param["typ"][0]
             if len(_param["typ"]) == 1
-            else "Union[{}]".format("|".join(_param["typ"]))
+            else "Union[{}]".format(",".join(_param["typ"]))
         )
     elif "$ref" in _param:
         _param["typ"] = transform_ref_fk_set(_param.pop("$ref"), fk)

diff --git a/cdd/tests/test_emit/test_emit_sqlalchemy_utils.py b/cdd/tests/test_emit/test_emit_sqlalchemy_utils.py
@@ -124,7 +124,7 @@ def test_update_args_infer_typ_sqlalchemy_when_simple_union(self) -> None:
         """Tests that SQLalchemy can infer the typ from a simple Union"""
         args = []
         update_args_infer_typ_sqlalchemy(
-            {"typ": "Union[string | Small]"}, args, "", False, {}
+            {"typ": "Union[string, Small]"}, args, "", False, {}
         )
         self.assertEqual(len(args), 1)
         run_ast_test(