From d4f8ccf3c4250e8861b23c2916eb27e78e032e6e Mon Sep 17 00:00:00 2001 From: Vlad Emelianov Date: Wed, 24 Apr 2024 23:38:54 +0300 Subject: [PATCH] Mark some output keys as optional --- mypy_boto3_builder/cli_parser.py | 3 +- mypy_boto3_builder/constants.py | 64 +----------------- mypy_boto3_builder/enums/product.py | 67 +++++++++++++++++++ .../enums/service_module_name.py | 4 +- .../generators/base_generator.py | 2 +- mypy_boto3_builder/main.py | 2 +- mypy_boto3_builder/parsers/shape_parser.py | 47 ++----------- mypy_boto3_builder/parsers/typed_dict_map.py | 47 +++++++++++++ .../type_annotations/type_typed_dict.py | 6 ++ .../type_annotations/test_type_typed_dict.py | 6 ++ 10 files changed, 139 insertions(+), 109 deletions(-) create mode 100644 mypy_boto3_builder/enums/product.py create mode 100644 mypy_boto3_builder/parsers/typed_dict_map.py diff --git a/mypy_boto3_builder/cli_parser.py b/mypy_boto3_builder/cli_parser.py index 3c72e06a..19bafe9d 100644 --- a/mypy_boto3_builder/cli_parser.py +++ b/mypy_boto3_builder/cli_parser.py @@ -8,7 +8,8 @@ from dataclasses import dataclass from pathlib import Path -from mypy_boto3_builder.constants import PROG_NAME, Product +from mypy_boto3_builder.constants import PROG_NAME +from mypy_boto3_builder.enums.product import Product from mypy_boto3_builder.service_name import ServiceName from mypy_boto3_builder.utils.version import get_builder_version diff --git a/mypy_boto3_builder/constants.py b/mypy_boto3_builder/constants.py index 453d664e..b58e0266 100644 --- a/mypy_boto3_builder/constants.py +++ b/mypy_boto3_builder/constants.py @@ -2,7 +2,6 @@ Constants and paths. """ -from enum import Enum from pathlib import Path # Random region to initialize services @@ -39,64 +38,5 @@ # universal mask for all resources ALL = "*" - -class ProductLibrary(Enum): - """ - Product library for Generator. - """ - - boto3 = "boto3" - aiobotocore = "aiobotocore" - aioboto3 = "aioboto3" - - -class ProductType(Enum): - """ - Product type for Generator. - """ - - stubs = "stubs" - service_stubs = "service_stubs" - docs = "docs" - - -class Product(Enum): - """ - Product choice for CLI. - """ - - boto3 = "boto3" - boto3_services = "boto3-services" - boto3_docs = "boto3-docs" - aiobotocore = "aiobotocore" - aiobotocore_services = "aiobotocore-services" - aiobotocore_docs = "aiobotocore-docs" - aioboto3 = "aioboto3" - aioboto3_docs = "aioboto3-docs" - - def __str__(self) -> str: - """ - Get string representation for debugging. - """ - return self.value - - def get_library(self) -> ProductLibrary: - """ - Get library name. - """ - for library in ProductLibrary: - if self.value.startswith(library.value): - return library - raise ValueError(f"No library found for {self.value}") - - def get_type(self) -> ProductType: - """ - Get product type. - """ - if "-" not in self.value: - return ProductType.stubs - if self.value.endswith("-services"): - return ProductType.service_stubs - if self.value.endswith("-docs"): - return ProductType.docs - raise ValueError(f"No type found for {self.value}") +# keys to mark as NotRequired for output TypeDicts +NOT_REQUIRED_OUTPUT_KEYS = ("NextToken", "Contents", "Item", "CommonPrefixes") diff --git a/mypy_boto3_builder/enums/product.py b/mypy_boto3_builder/enums/product.py new file mode 100644 index 00000000..ec6df51a --- /dev/null +++ b/mypy_boto3_builder/enums/product.py @@ -0,0 +1,67 @@ +""" +Product-related enums. +""" + +from enum import Enum + + +class ProductLibrary(Enum): + """ + Product library for Generator. + """ + + boto3 = "boto3" + aiobotocore = "aiobotocore" + aioboto3 = "aioboto3" + + +class ProductType(Enum): + """ + Product type for Generator. + """ + + stubs = "stubs" + service_stubs = "service_stubs" + docs = "docs" + + +class Product(Enum): + """ + Product choice for CLI. + """ + + boto3 = "boto3" + boto3_services = "boto3-services" + boto3_docs = "boto3-docs" + aiobotocore = "aiobotocore" + aiobotocore_services = "aiobotocore-services" + aiobotocore_docs = "aiobotocore-docs" + aioboto3 = "aioboto3" + aioboto3_docs = "aioboto3-docs" + + def __str__(self) -> str: + """ + Get string representation for debugging. + """ + return self.value + + def get_library(self) -> ProductLibrary: + """ + Get library name. + """ + for library in ProductLibrary: + if self.value.startswith(library.value): + return library + raise ValueError(f"No library found for {self.value}") + + def get_type(self) -> ProductType: + """ + Get product type. + """ + if "-" not in self.value: + return ProductType.stubs + if self.value.endswith("-services"): + return ProductType.service_stubs + if self.value.endswith("-docs"): + return ProductType.docs + raise ValueError(f"No type found for {self.value}") diff --git a/mypy_boto3_builder/enums/service_module_name.py b/mypy_boto3_builder/enums/service_module_name.py index 7696e2ad..bee43f8d 100644 --- a/mypy_boto3_builder/enums/service_module_name.py +++ b/mypy_boto3_builder/enums/service_module_name.py @@ -2,10 +2,10 @@ Enum for service modules. """ -import enum +from enum import Enum -class ServiceModuleName(enum.Enum): +class ServiceModuleName(Enum): """ Enum for service modules. """ diff --git a/mypy_boto3_builder/generators/base_generator.py b/mypy_boto3_builder/generators/base_generator.py index 01a9295c..9070668a 100644 --- a/mypy_boto3_builder/generators/base_generator.py +++ b/mypy_boto3_builder/generators/base_generator.py @@ -6,7 +6,7 @@ from collections.abc import Sequence from pathlib import Path -from mypy_boto3_builder.constants import ProductType +from mypy_boto3_builder.enums.product import ProductType from mypy_boto3_builder.logger import get_logger from mypy_boto3_builder.package_data import BasePackageData from mypy_boto3_builder.parsers.service_package_parser import ServicePackageParser diff --git a/mypy_boto3_builder/main.py b/mypy_boto3_builder/main.py index eba4940f..a750416a 100644 --- a/mypy_boto3_builder/main.py +++ b/mypy_boto3_builder/main.py @@ -9,7 +9,7 @@ from botocore.session import Session as BotocoreSession from mypy_boto3_builder.cli_parser import CLINamespace, parse_args -from mypy_boto3_builder.constants import Product, ProductLibrary +from mypy_boto3_builder.enums.product import Product, ProductLibrary from mypy_boto3_builder.generators.aioboto3_generator import AioBoto3Generator from mypy_boto3_builder.generators.aiobotocore_generator import AioBotocoreGenerator from mypy_boto3_builder.generators.base_generator import BaseGenerator diff --git a/mypy_boto3_builder/parsers/shape_parser.py b/mypy_boto3_builder/parsers/shape_parser.py index eca42d02..50380c74 100644 --- a/mypy_boto3_builder/parsers/shape_parser.py +++ b/mypy_boto3_builder/parsers/shape_parser.py @@ -3,7 +3,7 @@ """ import contextlib -from collections.abc import Iterable, Iterator, Sequence +from collections.abc import Iterable, Sequence from boto3.resources.model import Collection from boto3.session import Session @@ -21,6 +21,7 @@ ) from botocore.session import Session as BotocoreSession +from mypy_boto3_builder.constants import NOT_REQUIRED_OUTPUT_KEYS from mypy_boto3_builder.logger import get_logger from mypy_boto3_builder.parsers.shape_parser_types import ( ActionShape, @@ -30,6 +31,7 @@ ResourcesShape, WaitersShape, ) +from mypy_boto3_builder.parsers.typed_dict_map import TypedDictMap from mypy_boto3_builder.service_name import ServiceName from mypy_boto3_builder.structures.argument import Argument from mypy_boto3_builder.structures.method import Method @@ -61,7 +63,6 @@ ) from mypy_boto3_builder.utils.boto3_utils import get_botocore_session from mypy_boto3_builder.utils.strings import get_type_def_name -from mypy_boto3_builder.utils.type_def_sorter import TypeDefSorter class ShapeParserError(Exception): @@ -70,45 +71,6 @@ class ShapeParserError(Exception): """ -class TypedDictMap(dict[str, TypeTypedDict]): - """ - Wrapper for TypedDict maps. - """ - - def add(self, item: TypeTypedDict) -> None: - """ - Add new item. - """ - self[item.name] = item - - def iterate_pairs(self, name: str) -> Iterator[tuple[str, TypeTypedDict]]: - """ - Iterate over pairs mathed by real dict name. - """ - for key, value in list(self.items()): - if value.name == name: - yield key, value - - def rename(self, item: TypeTypedDict, new_name: str) -> None: - """ - Rename item and change mapping. - """ - for key, value in list(self.items()): - if value == item: - del self[key] - - item.name = new_name - self[new_name] = item - - def get_sorted_names(self) -> list[str]: - """ - Get real TypedDict names topologically sorted. - """ - sorted_values = TypeDefSorter(self.values()).sort() - allowed_names = {i.name for i in self.values()} - return [i.name for i in sorted_values if i.name in allowed_names] - - class ShapeParser: """ Parser for botocore shape files. @@ -455,7 +417,8 @@ def _parse_shape_structure( def _mark_typed_dict_as_total(self, typed_dict: TypeTypedDict) -> None: for attribute in typed_dict.children: - attribute.required = True + if attribute.name not in NOT_REQUIRED_OUTPUT_KEYS: + attribute.mark_as_required() def _add_response_metadata(self, typed_dict: TypeTypedDict) -> None: child_names = {i.name for i in typed_dict.children} diff --git a/mypy_boto3_builder/parsers/typed_dict_map.py b/mypy_boto3_builder/parsers/typed_dict_map.py new file mode 100644 index 00000000..2e6e70b1 --- /dev/null +++ b/mypy_boto3_builder/parsers/typed_dict_map.py @@ -0,0 +1,47 @@ +""" +Wrapper for TypedDict maps. +""" + +from collections.abc import Iterator + +from mypy_boto3_builder.type_annotations.type_typed_dict import TypeTypedDict +from mypy_boto3_builder.utils.type_def_sorter import TypeDefSorter + + +class TypedDictMap(dict[str, TypeTypedDict]): + """ + Wrapper for TypedDict maps. + """ + + def add(self, item: TypeTypedDict) -> None: + """ + Add new item. + """ + self[item.name] = item + + def iterate_pairs(self, name: str) -> Iterator[tuple[str, TypeTypedDict]]: + """ + Iterate over pairs mathed by real dict name. + """ + for key, value in list(self.items()): + if value.name == name: + yield key, value + + def rename(self, item: TypeTypedDict, new_name: str) -> None: + """ + Rename item and change mapping. + """ + for key, value in list(self.items()): + if value == item: + del self[key] + + item.name = new_name + self[new_name] = item + + def get_sorted_names(self) -> list[str]: + """ + Get real TypedDict names topologically sorted. + """ + sorted_values = TypeDefSorter(self.values()).sort() + allowed_names = {i.name for i in self.values()} + return [i.name for i in sorted_values if i.name in allowed_names] diff --git a/mypy_boto3_builder/type_annotations/type_typed_dict.py b/mypy_boto3_builder/type_annotations/type_typed_dict.py index 16f73fd5..05cf32e4 100644 --- a/mypy_boto3_builder/type_annotations/type_typed_dict.py +++ b/mypy_boto3_builder/type_annotations/type_typed_dict.py @@ -69,6 +69,12 @@ def is_required(self) -> bool: """ return self.required + def mark_as_required(self) -> None: + """ + Mark attribute as required. + """ + self.required = True + class TypeTypedDict(FakeAnnotation, TypeDefSortable): """ diff --git a/tests/type_annotations/test_type_typed_dict.py b/tests/type_annotations/test_type_typed_dict.py index ba001d61..96d34b16 100644 --- a/tests/type_annotations/test_type_typed_dict.py +++ b/tests/type_annotations/test_type_typed_dict.py @@ -16,6 +16,12 @@ def test_init(self) -> None: def test_render(self) -> None: assert self.result.render() == '"test": Dict[str, Any]' + def test_mark_as_required(self) -> None: + self.result.required = False + assert not self.result.is_required() + self.result.mark_as_required() + assert self.result.is_required() + class TestTypeTypedDict: result: TypeTypedDict