Use clashihg TypedDict logic for stubbed methods

Author: vemel

mypy_boto3_builder/parsers/client.py

@@ -18,6 +18,7 @@
 from mypy_boto3_builder.type_annotations.type import Type
 from mypy_boto3_builder.type_annotations.type_subscript import TypeSubscript
 from mypy_boto3_builder.type_maps.service_stub_map import get_stub_method_map
+from mypy_boto3_builder.utils.strings import get_class_prefix
 
 
 def parse_client(service_name: ServiceName, shape_parser: ShapeParser) -> Client:
@@ -44,6 +45,11 @@ def parse_client(service_name: ServiceName, shape_parser: ShapeParser) -> Client
     shape_method_map = shape_parser.get_client_method_map()
     stub_method_map = get_stub_method_map(service_name, parent_name)
     method_map = {**shape_method_map, **stub_method_map}
+    for method in stub_method_map.values():
+        shape_parser.create_request_type_annotation(
+            method=method,
+            name=f"{parent_name}{get_class_prefix(method.name)}Request",
+        )
 
     result.methods.append(result.get_exceptions_property())
     result.methods.extend(method_map.values())

mypy_boto3_builder/parsers/resource_parser.py

@@ -19,6 +19,7 @@
 from mypy_boto3_builder.structures.resource_record import ResourceRecord
 from mypy_boto3_builder.type_annotations.internal_import import InternalImport
 from mypy_boto3_builder.type_maps.service_stub_map import get_stub_method_map
+from mypy_boto3_builder.utils.strings import get_class_prefix
 
 
 class ResourceParser:
@@ -58,6 +59,11 @@ def resource_model(self) -> ResourceModel:
     def _parse_method_map(self) -> dict[str, Method]:
         shape_method_map = self.shape_parser.get_resource_method_map(self.name)
         stub_method_map = get_stub_method_map(self.service_name, self.name)
+        for method in stub_method_map.values():
+            self.shape_parser.create_request_type_annotation(
+                method=method,
+                name=f"{self.name}{get_class_prefix(method.name)}Request",
+            )
         return {**shape_method_map, **stub_method_map}
 
     def _parse_attributes(self, collections: Iterable[Collection]) -> list[Attribute]:

mypy_boto3_builder/parsers/service_resource_parser.py

@@ -16,6 +16,7 @@
 from mypy_boto3_builder.structures.service_resource import ServiceResource
 from mypy_boto3_builder.type_annotations.internal_import import InternalImport
 from mypy_boto3_builder.type_maps.service_stub_map import get_stub_method_map
+from mypy_boto3_builder.utils.strings import get_class_prefix
 
 
 class ServiceResourceParser(ResourceParser):
@@ -52,6 +53,11 @@ def _get_meta_attribute(self, name: str) -> Attribute:
     def _parse_method_map(self) -> dict[str, Method]:
         shape_method_map = self.shape_parser.get_service_resource_method_map()
         stub_method_map = get_stub_method_map(self.service_name, self.name)
+        for method in stub_method_map.values():
+            self.shape_parser.create_request_type_annotation(
+                method=method,
+                name=f"{self.name}{get_class_prefix(method.name)}Request",
+            )
         return {**shape_method_map, **stub_method_map}
 
     def _get_resource_model(self) -> ResourceModel:

mypy_boto3_builder/parsers/shape_parser.py

@@ -353,37 +353,38 @@ def get_client_method_map(self) -> dict[str, Method]:
                 ),
             )
             if operation_model.input_shape:
-                self._create_request_type_annotation(
+                self.create_request_type_annotation(
                     method=method,
-                    input_shape=operation_model.input_shape,
+                    name=operation_model.input_shape.name,
                     optional_postfix="Request",
                 )
             result[method.name] = method
         return result
 
-    def _create_request_type_annotation(
-        self, method: Method, input_shape: Shape, optional_postfix: str
+    def create_request_type_annotation(
+        self, method: Method, name: str, optional_postfix: str = "Extra"
     ) -> None:
+        """
+        Create a request type annotation for the given method and input shape.
+        """
         method.create_request_type_annotation(
-            self._get_non_clashing_typed_dict_name_for_shape(
-                input_shape,
-                optional_postfix="Request",
+            self._get_non_clashing_typed_dict_name(
+                name,
+                optional_postfix=optional_postfix,
             )
         )
         if method.has_request_type_annotation():
             self._request_typed_dict_map.add(method.request_type_annotation)
 
-    def _get_non_clashing_typed_dict_name_for_shape(
-        self, shape: Shape, optional_postfix: str
-    ) -> str:
-        temp_typed_dict = TypeTypedDict(self._get_typed_dict_name(shape), ())
-        return self._get_non_clashing_typed_dict_name(
+    def _get_non_clashing_typed_dict_name(self, name: str, optional_postfix: str) -> str:
+        temp_typed_dict = TypeTypedDict(get_type_def_name(name), ())
+        return self._get_non_clashing_typed_dict_name_for_existing(
             temp_typed_dict,
             optional_postfix=optional_postfix,
         )
 
     @staticmethod
-    def _get_typed_dict_name(shape: Shape, postfix: str = "") -> str:
+    def _get_typed_dict_name_for_shape(shape: Shape, postfix: str = "") -> str:
         return get_type_def_name(shape.name, postfix)
 
     @staticmethod
@@ -478,7 +479,9 @@ def _parse_shape_structure(
         typed_dict = TypeTypedDict(typed_dict_name)
 
         typed_dict_map = self._get_typed_dict_map(output=output, output_child=output_child)
-        resource_typed_dict_name = self._get_typed_dict_name(shape, postfix=self.resource_name)
+        resource_typed_dict_name = self._get_typed_dict_name_for_shape(
+            shape, postfix=self.resource_name
+        )
         found_typed_dict = typed_dict_map.get(typed_dict.name)
         found_resource_typed_dict = typed_dict_map.get(resource_typed_dict_name)
 
@@ -550,7 +553,7 @@ def _parse_shape_list(
 
     def _get_shape_type_name(self, shape: Shape) -> str:
         if isinstance(shape, StructureShape):
-            return self._get_typed_dict_name(shape)
+            return self._get_typed_dict_name_for_shape(shape)
 
         if isinstance(shape, StringShape):
             return shape.name
@@ -740,9 +743,9 @@ def get_paginate_method(self, paginator_name: str) -> Method:
             type_ignore="override",
         )
         if operation_model.input_shape is not None:
-            self._create_request_type_annotation(
+            self.create_request_type_annotation(
                 method=method,
-                input_shape=operation_model.input_shape,
+                name=operation_model.input_shape.name,
                 optional_postfix="Paginate",
             )
         return method
@@ -792,9 +795,9 @@ def get_wait_method(self, waiter_name: str) -> Method:
             type_ignore="override",
         )
         if operation_model.input_shape is not None:
-            self._create_request_type_annotation(
+            self.create_request_type_annotation(
                 method=method,
-                input_shape=operation_model.input_shape,
+                name=operation_model.input_shape.name,
                 optional_postfix="Wait",
             )
         return method
@@ -1123,9 +1126,9 @@ def _get_resource_method(self, action_name: str, action_shape: ActionShape) -> M
             ),
         )
         if operation_model.input_shape is not None:
-            self._create_request_type_annotation(
+            self.create_request_type_annotation(
                 method=method,
-                input_shape=operation_model.input_shape,
+                name=operation_model.input_shape.name,
                 optional_postfix=f"{self.resource_name}{action_name}",
             )
         return method
@@ -1233,7 +1236,7 @@ def _get_typed_dict(
                 return typed_dict_map[name]
         return None
 
-    def _get_non_clashing_typed_dict_name(
+    def _get_non_clashing_typed_dict_name_for_existing(
         self,
         typed_dict: TypeTypedDict,
         postfix: str = "",
@@ -1264,7 +1267,10 @@ def _get_non_clashing_typed_dict_name(
             f"Clashing typed dict name found: {new_typed_dict_name}",
             tags=new_typed_dict_name,
         )
-        return self._get_non_clashing_typed_dict_name(temp_typed_dict, postfix=optional_postfix)
+        return self._get_non_clashing_typed_dict_name_for_existing(
+            temp_typed_dict,
+            postfix=optional_postfix,
+        )
 
     def fix_typed_dict_names(self) -> None:
         """
@@ -1284,7 +1290,7 @@ def fix_typed_dict_names(self) -> None:
                     continue
 
                 old_typed_dict_name = output_typed_dict.name
-                new_typed_dict_name = self._get_non_clashing_typed_dict_name(
+                new_typed_dict_name = self._get_non_clashing_typed_dict_name_for_existing(
                     output_typed_dict,
                     "Output",
                 )
@@ -1318,7 +1324,7 @@ def fix_typed_dict_names(self) -> None:
                     continue
 
                 old_typed_dict_name = response_typed_dict.name
-                new_typed_dict_name = self._get_non_clashing_typed_dict_name(
+                new_typed_dict_name = self._get_non_clashing_typed_dict_name_for_existing(
                     response_typed_dict,
                     postfix="Response",
                 )
@@ -1346,7 +1352,9 @@ def convert_input_arguments_to_unions(
             if argument.type_annotation and is_type_parent(argument.type_annotation)
         )
         for input_typed_dict, output_typed_dict in self._fixed_typed_dict_map.items():
-            union_name = self._get_non_clashing_typed_dict_name(input_typed_dict, postfix="Union")
+            union_name = self._get_non_clashing_typed_dict_name_for_existing(
+                input_typed_dict, postfix="Union"
+            )
             union_type_annotation = TypeUnion(
                 name=union_name,
                 children=(input_typed_dict, output_typed_dict),

mypy_boto3_builder/type_maps/service_stub_map/__init__.py

@@ -11,7 +11,6 @@
 from mypy_boto3_builder.service_name import ServiceName, ServiceNameCatalog
 from mypy_boto3_builder.structures.method import Method
 from mypy_boto3_builder.type_maps.service_stub_map import dynamodb, ec2, rds, s3
-from mypy_boto3_builder.utils.strings import get_type_def_name
 
 ClassTypeMap = Mapping[str, Sequence[Method]]
 ServiceStubMap = Mapping[ServiceName, ClassTypeMap]
@@ -52,6 +51,4 @@ def get_stub_method_map(service_name: ServiceName, parent: str) -> dict[str, Met
     Get boto3 injected methods.
     """
     methods = SERVICE_STUB_MAP.get(service_name, {}).get(parent, [])
-    for method in methods:
-        method.create_request_type_annotation(get_type_def_name(parent, method.name, "Request"))
     return {method.name: method.copy() for method in methods}

scripts/integration.py

@@ -271,8 +271,8 @@ def parse_args() -> CLINamespace:
         "--product",
         nargs="*",
         choices=product_names,
-        default=[product_names[0]],
-        help=f"Product to test. Default: {product_names[0]}",
+        default=product_names,
+        help="Product to test. Default: all",
     )
     parser.add_argument(
         "--python",
@@ -488,8 +488,8 @@ def main() -> None:
         logger.error(error)
         sys.exit(1)
 
-    # if TEMP_PATH.exists():
-    #     shutil.rmtree(TEMP_PATH)
+    if TEMP_PATH.exists():
+        shutil.rmtree(TEMP_PATH)
 
 
 if __name__ == "__main__":