def _get_type_dict(self) -> FakeAnnotation: if not self.dict_items: return Type.DictStrAny first_key = self.dict_items[0]["key"] if first_key in SYNTAX_TYPE_MAP: result = TypeSubscript(Type.Dict) result.add_child(SYNTAX_TYPE_MAP[first_key]) result.add_child( TypeValue(self.service_name, self.prefix, self.dict_items[0]["value"]).get_type()) return result typed_dict_name = f"{self.prefix}TypeDef" shape_type_stub = get_shape_type_stub(self.service_name, typed_dict_name) if shape_type_stub: return shape_type_stub typed_dict = TypeTypedDict(typed_dict_name) for item in self.dict_items: key_name = self._parse_constant(item["key"]) prefix = f"{self.prefix}{key_name}" typed_dict.add_attribute( key_name, TypeValue(self.service_name, prefix, item["value"]).get_type(), required=False, ) return typed_dict
def _parse_shape_structure(self, shape: StructureShape) -> FakeAnnotation: if not shape.members.items(): return Type.DictStrAny required = shape.required_members typed_dict_name = f"{shape.name}TypeDef" shape_type_stub = get_shape_type_stub(self.service_name, typed_dict_name) if shape_type_stub: return shape_type_stub if typed_dict_name in self._typed_dict_map: return self._typed_dict_map[typed_dict_name] typed_dict = TypeTypedDict(typed_dict_name) self._typed_dict_map[typed_dict_name] = typed_dict for attr_name, attr_shape in shape.members.items(): typed_dict.add_attribute( attr_name, self._parse_shape(attr_shape), attr_name in required, ) if shape.name.endswith("Output"): typed_dict.add_attribute( "ResponseMetadata", self.response_metadata_typed_dict, False, ) return typed_dict
def _make_output_typed_dict(self, typed_dict: TypeTypedDict) -> None: for attribute in typed_dict.children: attribute.required = True child_names = {i.name for i in typed_dict.children} if "ResponseMetadata" not in child_names: typed_dict.add_attribute( "ResponseMetadata", self.response_metadata_typed_dict, True, )
def _parse_shape_structure( self, shape: StructureShape, output: bool = False, output_child: bool = False, is_streaming: bool = False, ) -> FakeAnnotation: if not shape.members.items(): return Type.DictStrAny if output_child else Type.MappingStrAny required = shape.required_members typed_dict_name = self._get_typed_dict_name(shape) shape_type_stub = get_shape_type_stub(self.service_name, typed_dict_name) if shape_type_stub: return shape_type_stub typed_dict = TypeTypedDict(typed_dict_name) if typed_dict.name in self._typed_dict_map: old_typed_dict = self._typed_dict_map[typed_dict.name] child_names = {i.name for i in old_typed_dict.children} if output and "ResponseMetadata" in child_names: return self._typed_dict_map[typed_dict.name] if not output and "ResponseMetadata" not in child_names: return self._typed_dict_map[typed_dict.name] if output: typed_dict.name = self._get_typed_dict_name( shape, postfix="ResponseMetadata") self.logger.debug( f"Marking {typed_dict.name} as ResponseMetadataTypeDef") else: old_typed_dict.name = self._get_typed_dict_name( shape, postfix="ResponseMetadata") self._typed_dict_map[old_typed_dict.name] = old_typed_dict self.logger.debug( f"Marking {old_typed_dict.name} as ResponseMetadataTypeDef" ) self._typed_dict_map[typed_dict.name] = typed_dict for attr_name, attr_shape in shape.members.items(): typed_dict.add_attribute( attr_name, self.parse_shape( attr_shape, output_child=output or output_child, is_streaming=is_streaming, ), attr_name in required, ) if output: self._make_output_typed_dict(typed_dict) return typed_dict
def _parse_shape_structure(self, shape: StructureShape) -> FakeAnnotation: if not shape.members.items(): return Type.DictStrAny required = shape.required_members typed_dict_name = f"{shape.name}TypeDef" if typed_dict_name in self._typed_dict_map: return self._typed_dict_map[typed_dict_name] typed_dict = TypeTypedDict(typed_dict_name) self._typed_dict_map[typed_dict_name] = typed_dict for attr_name, attr_shape in shape.members.items(): typed_dict.add_attribute( attr_name, self._parse_shape(attr_shape), attr_name in required, ) return typed_dict
def get_request_type_annotation(self, name: str) -> TypeTypedDict | None: """ Get TypedDict based on function arguments. """ result = TypeTypedDict(name) for argument in self.arguments: if argument.is_kwflag(): continue if not argument.type_annotation: continue result.add_attribute( argument.name, argument.type_annotation, required=argument.required, ) if not result.children: return None return result
class TestTypeTypedDict: result: TypeTypedDict def setup_method(self): self.result = TypeTypedDict( "MyDict", [ TypedDictAttribute("required", Type.bool, True), TypedDictAttribute("optional", Type.str, False), ], "documentation", ) def test_init(self) -> None: assert self.result.name == "MyDict" assert len(self.result.children) == 2 assert self.result.docstring == "documentation" assert self.result.requires_safe_render def test_get_attribute(self) -> None: assert self.result.get_attribute("required") == self.result.children[0] with pytest.raises(ValueError): self.result.get_attribute("non_existing") def test_render(self) -> None: assert self.result.render() == "MyDict" assert self.result.render("OtherDict") == "MyDict" assert self.result.render("MyDict") == '"MyDict"' self.result.stringify = True assert self.result.render() == '"MyDict"' self.result.replace_with_dict.add(self.result.name) assert self.result.render("MyDict") == "Dict[str, Any]" def test_get_import_record(self) -> None: assert self.result.get_import_record().render( ) == "from .type_defs import MyDict" def test_get_types(self) -> None: assert self.result.get_types() == {self.result} def test_add_attribute(self) -> None: self.result.add_attribute("third", Type.int, False) assert len(self.result.children) == 3 def test_is_dict(self) -> None: assert self.result.is_dict() def test_is_typed_dict(self) -> None: assert self.result.is_typed_dict() def test_render_class(self) -> None: assert self.result.render_class( ) == "class MyDict:\n required: bool\n optional: str" def test_has_optional(self) -> None: assert self.result.has_optional() assert not TypeTypedDict( "MyDict", [TypedDictAttribute("required", Type.bool, True)], ).has_optional() assert TypeTypedDict( "MyDict", [TypedDictAttribute("optional", Type.str, False)], ).has_optional() assert not TypeTypedDict("MyDict", []).has_optional() def test_has_required(self) -> None: assert self.result.has_required() assert TypeTypedDict( "MyDict", [TypedDictAttribute("required", Type.bool, True)], ).has_required() assert not TypeTypedDict( "MyDict", [TypedDictAttribute("optional", Type.str, False)], ).has_required() assert not TypeTypedDict("MyDict", []).has_required() def test_has_both(self) -> None: assert self.result.has_both() assert not TypeTypedDict( "MyDict", [TypedDictAttribute("required", Type.bool, True)], ).has_both() assert not TypeTypedDict( "MyDict", [TypedDictAttribute("optional", Type.str, False)], ).has_both() assert not TypeTypedDict("MyDict", []).has_both() def test_get_required(self) -> None: assert len(self.result.get_required()) == 1 assert self.result.get_required()[0].name == "required" def test_get_optional(self) -> None: assert len(self.result.get_optional()) == 1 assert self.result.get_optional()[0].name == "optional" def test_copy(self) -> None: assert self.result.copy().name == self.result.name def test_is_same(self) -> None: assert self.result.is_same( TypeTypedDict( "OtherDict", [ TypedDictAttribute("required", Type.bool, True), TypedDictAttribute("optional", Type.str, False), ], )) assert not self.result.is_same( TypeTypedDict( "OtherDict", [ TypedDictAttribute("required", Type.bool, True), TypedDictAttribute("optional", Type.float, False), ], )) def test_get_children_types(self) -> None: assert self.result.get_children_types() == {Type.str, Type.bool} def test_get_children_typed_dicts(self) -> None: assert len(self.result.get_children_typed_dicts()) == 0 def test_get_children_literals(self) -> None: assert len(self.result.get_children_literals()) == 0 assert len(self.result.get_children_literals([self.result])) == 0 def test_replace_self_references(self) -> None: self.result.replace_self_references()