Fixed a corner case when resolving required imports & applying types

saltudelft · May 12, 2021 · 99c758b · 99c758b
1 parent 9f0a842
commit 99c758b
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 4 deletions.
diff --git a/libsa4py/cst_transformers.py b/libsa4py/cst_transformers.py
@@ -1048,7 +1048,8 @@ def find_required_modules(all_types):
             for _, a_node in all_types:
                 m = match.findall(a_node.annotation, match.Attribute(value=match.DoNotCare(), attr=match.DoNotCare()))
                 if len(m) != 0:
-                    req_mod.add([n.value for n in match.findall(m[0], match.Name(value=match.DoNotCare()))][0])
+                    for i in m:
+                        req_mod.add([n.value for n in match.findall(i, match.Name(value=match.DoNotCare()))][0])
             return req_mod
 
         req_imports = []

diff --git a/tests/examples/type_apply_ex.json b/tests/examples/type_apply_ex.json
@@ -11,7 +11,8 @@
         "x": "int",
         "l": "typing.List[typing.Tuple[int, int]]",
         "c": "defaultdict",
-        "df": "List[pandas.arrays.PandasArray]"
+        "df": "pandas.DataFrame",
+        "dff": "typing.List[pandas.arrays.PandasArray]"
     },
     "mod_var_occur": {
         "x": [

diff --git a/tests/test_type_apply.py b/tests/test_type_apply.py
@@ -8,7 +8,8 @@
 x: int = 12
 l = [(1, 2)]
 c = defaultdict(int)
-df = pd.Dataframe([1,2])
+df = pd.DataFrame([2, 3])
+dff = pd.DataFrame([1,2])
 class Foo:
     foo_v: str = 'Hello, Foo!'
     class Delta:
@@ -38,7 +39,8 @@ def Bar(x=['apple', 'orange']):
 x: int = 12
 l: typing.List[typing.Tuple[int, int]] = [(1, 2)]
 c: defaultdict = defaultdict(int)
-df: List[pandas.arrays.PandasArray] = pd.Dataframe([1,2])
+df: pandas.DataFrame = pd.DataFrame([2, 3])
+dff: typing.List[pandas.arrays.PandasArray] = pd.DataFrame([1,2])
 class Foo:
     foo_v: str = 'Hello, Foo!'
     class Delta: