Add --use-all-exports option to generate __all__ in __init__.py (#2586)

* Add --use-all-exports option to generate __all__ declaration in __init__.py

* Refactor __init__.py and base.py to implement --use-all-exports option for dynamic __all__ declaration generation

* Refactor base.py to streamline __all__ declaration generation for __init__.py with --use-all-exports option

Author: koxudaxi

README.md

@@ -490,6 +490,8 @@ Model customization:
                         target python version
   --treat-dot-as-module
                         treat dotted module names as modules
+  --use-all-exports     Generate __all__ = [...] in __init__.py to export all defined models
+                        and types
   --use-exact-imports   import exact types instead of modules, for example: "from .foo
                         import Bar" instead of "from . import foo" with "foo.Bar"
   --use-pendulum        use pendulum instead of datetime

docs/index.md

@@ -482,6 +482,8 @@ Model customization:
                         target python version
   --treat-dot-as-module
                         treat dotted module names as modules
+  --use-all-exports     Generate __all__ = [...] in __init__.py to export all defined models
+                        and types
   --use-exact-imports   import exact types instead of modules, for example: "from .foo
                         import Bar" instead of "from . import foo" with "foo.Bar"
   --use-pendulum        use pendulum instead of datetime

src/datamodel_code_generator/__init__.py

@@ -375,6 +375,7 @@ def generate(  # noqa: PLR0912, PLR0913, PLR0914, PLR0915
     dataclass_arguments: DataclassArguments | None = None,
     disable_future_imports: bool = False,
     type_mappings: list[str] | None = None,
+    use_all_exports: bool = False,
 ) -> None:
     """Generate Python data models from schema definitions or structured data.
 
@@ -612,7 +613,7 @@ def get_header_and_first_line(csv_file: IO[str]) -> dict[str, Any]:
     )
 
     with chdir(output):
-        results = parser.parse(disable_future_imports=disable_future_imports)
+        results = parser.parse(disable_future_imports=disable_future_imports, use_all_exports=use_all_exports)
     if not input_filename:  # pragma: no cover
         if isinstance(input_, str):
             input_filename = "<stdin>"
@@ -628,7 +629,7 @@ def get_header_and_first_line(csv_file: IO[str]) -> dict[str, Any]:
         msg = "Models not found in the input data"
         raise Error(msg)
     if isinstance(results, str):
-        modules = {output: (results, input_filename)}
+        modules: dict[Path | None, tuple[str, str | None]] = {output: (results, input_filename)}
     else:
         if output is None:
             msg = "Modular references require an output directory"

src/datamodel_code_generator/__main__.py

@@ -425,6 +425,7 @@ def validate_root(cls, values: dict[str, Any]) -> dict[str, Any]:  # noqa: N805
     parent_scoped_naming: bool = False
     disable_future_imports: bool = False
     type_mappings: Optional[list[str]] = None  # noqa: UP045
+    use_all_exports: bool = False
 
     def merge_args(self, args: Namespace) -> None:
         """Merge command-line arguments into config."""
@@ -826,6 +827,7 @@ def main(args: Sequence[str] | None = None) -> Exit:  # noqa: PLR0911, PLR0912,
             dataclass_arguments=config.dataclass_arguments,
             disable_future_imports=config.disable_future_imports,
             type_mappings=config.type_mappings,
+            use_all_exports=config.use_all_exports,
         )
     except InvalidClassNameError as e:
         print(f"{e} You have to set `--class-name` option", file=sys.stderr)  # noqa: T201

src/datamodel_code_generator/arguments.py

@@ -293,6 +293,12 @@ def start_section(self, heading: str | None) -> None:
     action="store_true",
     default=None,
 )
+model_options.add_argument(
+    "--use-all-exports",
+    help="Generate __all__ = [...] in __init__.py to export all defined models and types",
+    action="store_true",
+    default=None,
+)
 
 # ======================================================================================
 # Typing options for generated models

src/datamodel_code_generator/parser/base.py

@@ -1568,6 +1568,7 @@ def parse(  # noqa: PLR0912, PLR0914, PLR0915
         format_: bool | None = True,  # noqa: FBT001, FBT002
         settings_path: Path | None = None,
         disable_future_imports: bool = False,  # noqa: FBT001, FBT002
+        use_all_exports: bool = False,  # noqa: FBT001, FBT002
     ) -> str | dict[tuple[str, ...], Result]:
         """Parse schema and generate code, returning single file or module dict."""
         self.parse_raw()
@@ -1709,10 +1710,20 @@ class Processed(NamedTuple):
 
         for module, models, init, imports, scoped_model_resolver in processed_models:  # noqa: B007
             result: list[str] = []
+            export_names: list[str] = []
             if models:
                 if with_import:
                     result += [str(self.imports), str(imports), "\n"]
 
+                if use_all_exports and module[-1] == "__init__.py":
+                    export_names = [
+                        m.reference.short_name
+                        for m in models
+                        if m.reference and not m.reference.short_name.startswith("_")
+                    ]
+                    all_items = ",\n    ".join(f'"{name}"' for name in export_names)
+                    result += [f"__all__ = [\n    {all_items},\n]\n"]
+
                 self.__update_type_aliases(models)
                 code = dump_templates(models)
                 result += [code]

tests/data/custom_file_header_docstring.txt

@@ -0,0 +1,5 @@
+"""
+Custom module docstring header.
+
+This is a multi-line docstring used for testing.
+"""

tests/data/expected/main/openapi/modular_all_exports/__init__.py

@@ -0,0 +1,47 @@
+# generated by datamodel-codegen:
+#   filename:  modular.yaml
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel
+
+from . import foo as foo_1
+from . import models
+from .nested import foo as foo_2
+
+__all__ = [
+    "OptionalModel",
+    "Id",
+    "Error",
+    "Result",
+    "Source",
+    "DifferentTea",
+]
+
+
+class OptionalModel(BaseModel):
+    __root__: str
+
+
+class Id(BaseModel):
+    __root__: str
+
+
+class Error(BaseModel):
+    code: int
+    message: str
+
+
+class Result(BaseModel):
+    event: Optional[models.Event] = None
+
+
+class Source(BaseModel):
+    country: Optional[str] = None
+
+
+class DifferentTea(BaseModel):
+    foo: Optional[foo_1.Tea] = None
+    nested: Optional[foo_2.Tea] = None

tests/data/expected/main/openapi/modular_all_exports/bar.py

@@ -0,0 +1,10 @@
+# generated by datamodel-codegen:
+#   filename:  modular.yaml
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+
+class FieldModel(BaseModel):
+    __root__: str = Field(..., example='green')

tests/data/expected/main/openapi/modular_all_exports/collections.py

@@ -0,0 +1,50 @@
+# generated by datamodel-codegen:
+#   filename:  modular.yaml
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import List, Optional
+
+from pydantic import AnyUrl, BaseModel, Field
+
+from . import models
+
+
+class Pets(BaseModel):
+    __root__: List[models.Pet]
+
+
+class Users(BaseModel):
+    __root__: List[models.User]
+
+
+class Rules(BaseModel):
+    __root__: List[str]
+
+
+class Stage(Enum):
+    test = 'test'
+    dev = 'dev'
+    stg = 'stg'
+    prod = 'prod'
+
+
+class Api(BaseModel):
+    apiKey: Optional[str] = Field(
+        None, description='To be used as a dataset parameter value'
+    )
+    apiVersionNumber: Optional[str] = Field(
+        None, description='To be used as a version parameter value'
+    )
+    apiUrl: Optional[AnyUrl] = Field(
+        None, description="The URL describing the dataset's fields"
+    )
+    apiDocumentationUrl: Optional[AnyUrl] = Field(
+        None, description='A URL to the API console for each API'
+    )
+    stage: Optional[Stage] = None
+
+
+class Apis(BaseModel):
+    __root__: List[Api]