Mark some output keys as optional

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
 

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")

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}")

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.
     """

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

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

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}

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]

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):
     """

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