def generate_additional_get_all_method()

in src/sagemaker_core/tools/resources_codegen.py [0:0]

• 80 lines of code
• 9 McCabe index (conditional complexity)

    def generate_additional_get_all_method(self, method: Method, resource_attributes: list):
        """Auto-Generate methods that return a list of objects.

        Args:
            resource_name (str): The resource name.

        Returns:
            str: The formatted method code.
        """
        # TODO: merge this with generate_get_all_method
        operation_metadata = self.operations[method.operation_name]
        operation_input_shape_name = operation_metadata["input"]["shape"]
        exclude_list = ["next_token", "max_results"]
        if method.method_type == MethodType.CLASS.value:
            decorator = "@classmethod"
            method_args = add_indent("cls,\n", 4)
            method_args += self._generate_method_args(operation_input_shape_name, exclude_list)
            operation_input_args = self._generate_operation_input_args_updated(
                operation_metadata, True, resource_attributes, exclude_list
            )
            exclude_resource_attrs = None
        else:
            decorator = ""
            method_args = add_indent("self,\n", 4)
            method_args += self._generate_method_args(
                operation_input_shape_name, exclude_list + resource_attributes
            )
            operation_input_args = self._generate_operation_input_args_updated(
                operation_metadata, False, resource_attributes, exclude_list
            )
            exclude_resource_attrs = resource_attributes
        method_args += add_indent("session: Optional[Session] = None,\n", 4)
        method_args += add_indent("region: Optional[str] = None,", 4)

        iterator_return_type = method.return_type
        if method.resource_name != method.return_type and method.return_type in self.shapes:
            iterator_return_type = f"shapes.{method.return_type}"

        if method.return_type == method.resource_name:
            method_return_type = f'ResourceIterator["{method.resource_name}"]'
        else:
            method_return_type = f"ResourceIterator[{iterator_return_type}]"
        return_string = f"Returns:\n" f"    Iterator for listed {method.return_type}.\n"

        get_list_operation_output_shape = operation_metadata["output"]["shape"]
        list_operation_output_members = self.shapes[get_list_operation_output_shape]["members"]

        filtered_list_operation_output_members = next(
            {key: value}
            for key, value in list_operation_output_members.items()
            if key != "NextToken"
        )
        summaries_key = next(iter(filtered_list_operation_output_members))
        summaries_shape_name = filtered_list_operation_output_members[summaries_key]["shape"]
        summary_name = self.shapes[summaries_shape_name]["member"]["shape"]

        list_method = convert_to_snake_case(method.operation_name)

        # TODO: add rules for custom key mapping and list methods with no args
        resource_iterator_args_list = [
            "client=client",
            f"list_method='{list_method}'",
            f"summaries_key='{summaries_key}'",
            f"summary_name='{summary_name}'",
            f"resource_cls={iterator_return_type}",
            "list_method_kwargs=operation_input_args",
        ]

        resource_iterator_args = ",\n".join(resource_iterator_args_list)
        resource_iterator_args = add_indent(resource_iterator_args, 8)
        serialize_operation_input = SERIALIZE_INPUT_TEMPLATE.format(
            operation_input_args=operation_input_args
        )
        initialize_client = INITIALIZE_CLIENT_TEMPLATE.format(service_name=method.service_name)
        deserialize_response = RETURN_ITERATOR_TEMPLATE.format(
            resource_iterator_args=resource_iterator_args
        )

        # generate docstring
        docstring = self._generate_docstring(
            title=method.docstring_title,
            operation_name=method.operation_name,
            resource_name=method.resource_name,
            operation_input_shape_name=operation_input_shape_name,
            include_session_region=True,
            return_string=return_string,
            exclude_resource_attrs=exclude_resource_attrs,
        )

        return GENERIC_METHOD_TEMPLATE.format(
            docstring=docstring,
            decorator=decorator,
            method_name=method.method_name,
            method_args=method_args,
            return_type=method_return_type,
            serialize_operation_input=serialize_operation_input,
            initialize_client=initialize_client,
            call_operation_api="",
            deserialize_response=deserialize_response,
        )