diff --git a/libsa4py/cst_transformers.py b/libsa4py/cst_transformers.py index a95fee1..efef2dd 100644 --- 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 index 2d259da..9de3a46 100644 --- 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 index 7be976b..77993cc 100644 --- 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: