def generate_wait_method()

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

• 42 lines of code
• 7 McCabe index (conditional complexity)

    def generate_wait_method(self, resource_name: str) -> str:
        """Auto-Generate WAIT Method for a waitable resource.

        Args:
            resource_name (str): The resource name.

        Returns:
            str: The formatted Wait Method template.
        """
        resource_status_chain, resource_states = (
            self.resources_extractor.get_status_chain_and_states(resource_name)
        )

        # Get terminal states for resource
        terminal_resource_states = []
        for state in resource_states:
            # Handles when a resource has terminal states like UpdateCompleted, CreateFailed, etc.
            # Checking lower because case is not consistent accross resources (ie, COMPLETED vs Completed)
            if any(terminal_state.lower() in state.lower() for terminal_state in TERMINAL_STATES):
                terminal_resource_states.append(state)

        # Get resource status key path
        status_key_path = ""
        for member in resource_status_chain:
            status_key_path += f'.{convert_to_snake_case(member["name"])}'

        failure_reason = self._get_failure_reason_ref(resource_name)
        formatted_failed_block = FAILED_STATUS_ERROR_TEMPLATE.format(
            resource_name=resource_name, reason=failure_reason
        )
        formatted_failed_block = add_indent(formatted_failed_block, 16)

        logs_arg = ""
        logs_arg_doc = ""
        init_wait_logs = ""
        print_wait_logs = ""
        if resource_name in RESOURCE_WITH_LOGS:
            logs_arg = "logs: Optional[bool] = False,"
            logs_arg_doc = "logs: Whether to print logs while waiting.\n"

            instance_count = self._get_instance_count_ref(resource_name)
            init_wait_logs = add_indent(
                INIT_WAIT_LOGS_TEMPLATE.format(
                    get_instance_count=instance_count,
                    job_type=resource_name,
                )
            )
            print_wait_logs = add_indent(PRINT_WAIT_LOGS, 12)

        formatted_method = WAIT_METHOD_TEMPLATE.format(
            terminal_resource_states=terminal_resource_states,
            status_key_path=status_key_path,
            failed_error_block=formatted_failed_block,
            resource_name=resource_name,
            logs_arg=logs_arg,
            logs_arg_doc=logs_arg_doc,
            init_wait_logs=init_wait_logs,
            print_wait_logs=print_wait_logs,
        )
        return formatted_method