From b53e6709592adcee4b581c80babb562a771d4d60 Mon Sep 17 00:00:00 2001 From: Mufaddal Rohawala Date: Tue, 26 Nov 2024 18:20:06 -0800 Subject: [PATCH 1/2] Add support for map-in-list, list-in-list structures --- .../main/code_injection/codec.py | 18 ++++++++++++++++-- src/sagemaker_core/tools/method.py | 5 +++-- src/sagemaker_core/tools/resources_codegen.py | 9 +++++++-- src/sagemaker_core/tools/shapes_extractor.py | 19 ++++++++++++------- 4 files changed, 38 insertions(+), 13 deletions(-) diff --git a/src/sagemaker_core/main/code_injection/codec.py b/src/sagemaker_core/main/code_injection/codec.py index bb0fd365..286f7a16 100644 --- a/src/sagemaker_core/main/code_injection/codec.py +++ b/src/sagemaker_core/main/code_injection/codec.py @@ -115,16 +115,30 @@ def _evaluate_list_type(raw_list, shape) -> list: """ _shape_member_type = shape["member_type"] _shape_member_shape = shape["member_shape"] + _evaluated_list = [] if _shape_member_type in BASIC_TYPES: # if basic types directly assign list value. _evaluated_list = raw_list elif _shape_member_type == STRUCTURE_TYPE: - # if structure type transform each list item and assign value. - _evaluated_list = [] + # if structure type, transform each list item and assign value. # traverse through response list and evaluate item for item in raw_list: _evaluated_item = transform(item, _shape_member_shape) _evaluated_list.append(_evaluated_item) + elif _shape_member_type == LIST_TYPE: + # if list type, transform each list item and assign value. + # traverse through response list and evaluate item + for item in raw_list: + _list_type_shape = SHAPE_DAG[_shape_member_shape] + _evaluated_item = _evaluate_list_type(item, _list_type_shape) + _evaluated_list.append(_evaluated_item) + elif _shape_member_type == MAP_TYPE: + # if structure type, transform each list item and assign value. + # traverse through response list and evaluate item + for item in raw_list: + _map_type_shape = SHAPE_DAG[_shape_member_shape] + _evaluated_item = _evaluate_map_type(item, _map_type_shape) + _evaluated_list.append(_evaluated_item) else: raise ValueError( f"Unhandled List member type " diff --git a/src/sagemaker_core/tools/method.py b/src/sagemaker_core/tools/method.py index cc19b19b..18f075fb 100644 --- a/src/sagemaker_core/tools/method.py +++ b/src/sagemaker_core/tools/method.py @@ -26,7 +26,8 @@ def __init__(self, **kwargs): self.__dict__.update(kwargs) def get_docstring_title(self, operation): - title = remove_html_tags(operation["documentation"]) - self.docstring_title = title.split(".")[0] + "." + documentation = operation.get("documentation") + title = remove_html_tags(documentation) if documentation else None + self.docstring_title = title.split(".")[0] + "." if title else None # TODO: add some templates for common methods diff --git a/src/sagemaker_core/tools/resources_codegen.py b/src/sagemaker_core/tools/resources_codegen.py index 97cc253a..dca20be3 100644 --- a/src/sagemaker_core/tools/resources_codegen.py +++ b/src/sagemaker_core/tools/resources_codegen.py @@ -931,7 +931,9 @@ def _fetch_shape_errors_and_doc_strings(self, operation): for e in errors: error_shape = e["shape"] error_shape_dict = self.shapes[error_shape] - error_shape_documentation = error_shape_dict.get("documentation").strip() + error_shape_documentation = error_shape_dict.get("documentation") + if error_shape_documentation: + error_shape_documentation.strip() shape_errors_and_docstrings[error_shape] = error_shape_documentation sorted_keys = sorted(shape_errors_and_docstrings.keys()) return {key: shape_errors_and_docstrings[key] for key in sorted_keys} @@ -939,7 +941,10 @@ def _fetch_shape_errors_and_doc_strings(self, operation): def _exception_docstring(self, operation: str) -> str: _docstring = RESOURCE_METHOD_EXCEPTION_DOCSTRING for error, documentaion in self._fetch_shape_errors_and_doc_strings(operation).items(): - _docstring += f"\n {error}: {remove_html_tags(documentaion).strip()}" + if documentaion: + _docstring += f"\n {error}: {remove_html_tags(documentaion).strip()}" + else: + _docstring += f"\n {error}" return _docstring def _generate_docstring( diff --git a/src/sagemaker_core/tools/shapes_extractor.py b/src/sagemaker_core/tools/shapes_extractor.py index 0aefce22..7ca2a8d0 100644 --- a/src/sagemaker_core/tools/shapes_extractor.py +++ b/src/sagemaker_core/tools/shapes_extractor.py @@ -119,19 +119,24 @@ def get_shapes_dag(self): def _evaluate_list_type(self, member_shape): list_shape_name = member_shape["member"]["shape"] - list_shape_type = self.combined_shapes[list_shape_name]["type"] - if list_shape_type in ["list", "map"]: - raise Exception( - "Unhandled list shape key type encountered, needs extra logic to handle this" - ) - if list_shape_type == "structure": + list_shape_member = self.combined_shapes[list_shape_name] + list_shape_type = list_shape_member["type"] + if list_shape_type == "list": + member_type = f"List[{self._evaluate_list_type(list_shape_member)}]" + elif list_shape_type == "map": + member_type = f"List[{self._evaluate_map_type(list_shape_member)}]" + elif list_shape_type == "structure": # handling an edge case of nested structure if list_shape_name == "SearchExpression": member_type = f"List['{list_shape_name}']" else: member_type = f"List[{list_shape_name}]" - else: + elif list_shape_type in BASIC_JSON_TYPES_TO_PYTHON_TYPES.keys(): member_type = f"List[{BASIC_JSON_TYPES_TO_PYTHON_TYPES[list_shape_type]}]" + else: + raise Exception( + f"Unhandled list shape key type {list_shape_type} for Shape: {list_shape_name} encountered, needs extra logic to handle this" + ) return member_type def _evaluate_map_type(self, member_shape): From 685104e9e532a8c72e3eab639941334b3413f155 Mon Sep 17 00:00:00 2001 From: Erick Benitez-Ramos Date: Wed, 4 Dec 2024 12:55:51 -0800 Subject: [PATCH 2/2] update boto3 version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 8b199918..6fe874a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,7 +12,7 @@ authors = [ readme = "README.rst" dependencies = [ # Add your dependencies here (Include lower and upper bounds as applicable) - "boto3>=1.34.0,<2.0.0", + "boto3>=1.35.75,<2.0.0", "pydantic>=2.0.0,<3.0.0", "PyYAML>=6.0, <7.0", "jsonschema<5.0.0",