Merge pull request #2 from ag2ai/Bugfixing

Bugfixing

Author: sternakt

fastapi_code_generator/__init__.py

@@ -0,0 +1,3 @@
+from .patches import patch_parse
+
+patch_parse()

fastapi_code_generator/parser.py

@@ -93,16 +93,43 @@ class Argument(CachedPropertyModel):
     type_hint: UsefulStr
     default: Optional[UsefulStr] = None
     default_value: Optional[UsefulStr] = None
+    field: Union[DataModelField, list[DataModelField], None] = None
     required: bool
 
     def __str__(self) -> str:
         return self.argument
 
-    @cached_property
+    @property
     def argument(self) -> str:
+        if self.field is None:
+            type_hint = self.type_hint
+        else:
+            type_hint = (
+                UsefulStr(self.field.type_hint)
+                if not isinstance(self.field, list)
+                else UsefulStr(
+                    f"Union[{', '.join(field.type_hint for field in self.field)}]"
+                )
+            )
+        if self.default is None and self.required:
+            return f'{self.name}: {type_hint}'
+        return f'{self.name}: {type_hint} = {self.default}'
+
+    @property
+    def snakecase(self) -> str:
+        if self.field is None:
+            type_hint = self.type_hint
+        else:
+            type_hint = (
+                UsefulStr(self.field.type_hint)
+                if not isinstance(self.field, list)
+                else UsefulStr(
+                    f"Union[{', '.join(field.type_hint for field in self.field)}]"
+                )
+            )
         if self.default is None and self.required:
-            return f'{self.name}: {self.type_hint}'
-        return f'{self.name}: {self.type_hint} = {self.default}'
+            return f'{stringcase.snakecase(self.name)}: {type_hint}'
+        return f'{stringcase.snakecase(self.name)}: {type_hint} = {self.default}'
 
 
 class Operation(CachedPropertyModel):
@@ -114,16 +141,39 @@ class Operation(CachedPropertyModel):
     parameters: List[Dict[str, Any]] = []
     responses: Dict[UsefulStr, Any] = {}
     deprecated: bool = False
-    imports: List[Import] = []
     security: Optional[List[Dict[str, List[str]]]] = None
     tags: Optional[List[str]] = []
-    arguments: str = ''
-    snake_case_arguments: str = ''
     request: Optional[Argument] = None
     response: str = ''
     additional_responses: Dict[Union[str, int], Dict[str, str]] = {}
     return_type: str = ''
     callbacks: Dict[UsefulStr, List["Operation"]] = {}
+    arguments_list: List[Argument] = []
+
+    @classmethod
+    def merge_arguments_with_union(cls, arguments: List[Argument]) -> List[Argument]:
+        grouped_arguments: DefaultDict[str, List[Argument]] = DefaultDict(list)
+        for argument in arguments:
+            grouped_arguments[argument.name].append(argument)
+
+        merged_arguments = []
+        for argument_list in grouped_arguments.values():
+            if len(argument_list) == 1:
+                merged_arguments.append(argument_list[0])
+            else:
+                argument = argument_list[0]
+                fields = [
+                    item
+                    for arg in argument_list
+                    if arg.field is not None
+                    for item in (
+                        arg.field if isinstance(arg.field, list) else [arg.field]
+                    )
+                    if item is not None
+                ]
+                argument.field = fields
+                merged_arguments.append(argument)
+        return merged_arguments
 
     @cached_property
     def type(self) -> UsefulStr:
@@ -132,6 +182,27 @@ def type(self) -> UsefulStr:
         """
         return self.method
 
+    @property
+    def arguments(self) -> str:
+        sorted_arguments = Operation.merge_arguments_with_union(self.arguments_list)
+        return ", ".join(argument.argument for argument in sorted_arguments)
+
+    @property
+    def snake_case_arguments(self) -> str:
+        sorted_arguments = Operation.merge_arguments_with_union(self.arguments_list)
+        return ", ".join(argument.snakecase for argument in sorted_arguments)
+
+    @property
+    def imports(self) -> Imports:
+        imports = Imports()
+        for argument in self.arguments_list:
+            if isinstance(argument.field, list):
+                for field in argument.field:
+                    imports.append(field.data_type.import_)
+            elif argument.field:
+                imports.append(argument.field.data_type.import_)
+        return imports
+
     @cached_property
     def root_path(self) -> UsefulStr:
         paths = self.path.split("/")
@@ -314,6 +385,7 @@ def get_parameter_type(
             default=default,  # type: ignore
             default_value=schema.default,
             required=field.required,
+            field=field,
         )
 
     def get_arguments(self, snake_case: bool, path: List[str]) -> str:
@@ -347,6 +419,10 @@ def get_argument_list(self, snake_case: bool, path: List[str]) -> List[Argument]
                 or argument.type_hint.startswith('Optional[')
             )
 
+        # check if there are duplicate argument.name
+        argument_names = [argument.name for argument in arguments]
+        if len(argument_names) != len(set(argument_names)):
+            self.imports_for_fastapi.append(Import(from_='typing', import_="Union"))
         return arguments
 
     def parse_request_body(
@@ -466,10 +542,7 @@ def parse_operation(
         resolved_path = self.model_resolver.resolve_ref(path)
         path_name, method = path[-2:]
 
-        self._temporary_operation['arguments'] = self.get_arguments(
-            snake_case=False, path=path
-        )
-        self._temporary_operation['snake_case_arguments'] = self.get_arguments(
+        self._temporary_operation['arguments_list'] = self.get_argument_list(
             snake_case=True, path=path
         )
         main_operation = self._temporary_operation
@@ -499,11 +572,8 @@ def parse_operation(
                         self._temporary_operation = {'_parameters': []}
                         cb_path = path + ['callbacks', key, route, method]
                         super().parse_operation(cb_op, cb_path)
-                        self._temporary_operation['arguments'] = self.get_arguments(
-                            snake_case=False, path=cb_path
-                        )
-                        self._temporary_operation['snake_case_arguments'] = (
-                            self.get_arguments(snake_case=True, path=cb_path)
+                        self._temporary_operation['arguments_list'] = (
+                            self.get_argument_list(snake_case=True, path=cb_path)
                         )
 
                         callbacks[key].append(
@@ -527,13 +597,16 @@ def _collapse_root_model(self, data_type: DataType) -> DataType:
         reference = data_type.reference
         import functools
 
-        if not (
-            reference
-            and (
-                len(reference.children) == 1
-                or functools.reduce(lambda a, b: a == b, reference.children)
-            )
-        ):
+        try:
+            if not (
+                reference
+                and (
+                    len(reference.children) == 0
+                    or functools.reduce(lambda a, b: a == b, reference.children)
+                )
+            ):
+                return data_type
+        except RecursionError:
             return data_type
         source = reference.source
         if not isinstance(source, CustomRootType):