Replace List with list

Author: dgarros

infrahub_sdk/graphql/plugin.py

@@ -25,7 +25,7 @@ def replace_base_model_import(cls, module: ast.Module) -> ast.Module:
         if base_model_index == -1:
             raise ValueError("BaseModel not found in module")
         module.body[base_model_index] = ast.ImportFrom(
-            module="infrahub_sdk.graphql", names=[ast.alias(name=GraphQLReturnTypeModel.__name__)]
+            module="infrahub_sdk.graphql", names=[ast.alias(name=GraphQLReturnTypeModel.__name__)], level=2
         )
         return module
 
@@ -41,9 +41,34 @@ def replace_base_model_class(module: ast.Module) -> ast.Module:
                     base.id = GraphQLReturnTypeModel.__name__
         return module
 
-    def insert_future_annotation(self, module: ast.Module) -> ast.Module:
+    @staticmethod
+    def insert_future_annotation(module: ast.Module) -> ast.Module:
         """Insert the future annotation at the beginning of the module."""
-        module.body.insert(0, ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")]))
+        module.body.insert(0, ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0))
+        return module
+
+    @classmethod
+    def replace_list_in_subscript(cls, subscript: ast.Subscript) -> ast.Subscript:
+        if isinstance(subscript.value, ast.Name) and subscript.value.id == "List":
+            subscript.value.id = "list"
+        if isinstance(subscript.slice, ast.Subscript):
+            subscript.slice = cls.replace_list_in_subscript(subscript.slice)
+
+        return subscript
+
+    @classmethod
+    def replace_list_annotations(cls, module: ast.Module) -> ast.Module:
+        for item in module.body:
+            if not isinstance(item, ast.ClassDef):
+                continue
+
+            # replace List with list in the annotations when list is used as a type
+            for class_item in item.body:
+                if not isinstance(class_item, ast.AnnAssign):
+                    continue
+                if isinstance(class_item.annotation, ast.Subscript):
+                    class_item.annotation = cls.replace_list_in_subscript(class_item.annotation)
+
         return module
 
     def generate_result_types_module(
@@ -55,4 +80,6 @@ def generate_result_types_module(
         module = self.replace_base_model_import(module)
         module = self.replace_base_model_class(module)
 
+        module = self.replace_list_annotations(module)
+
         return module