Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: new cli option --use-exact-imports #1983

Merged
merged 7 commits into from
Jun 18, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -457,6 +457,9 @@ Model customization:
--use-schema-description
Use schema description to populate class docstring
--use-title-as-name use titles as class names of models
--use-exact-imports Import exact types instead of modules, for example:
`from .foo import Bar` instead of
`from . import foo` with `foo.Bar`

Template customization:
--aliases ALIASES Alias mapping file
Expand Down
2 changes: 2 additions & 0 deletions datamodel_code_generator/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -301,6 +301,7 @@ def generate(
custom_formatters_kwargs: Optional[Dict[str, Any]] = None,
use_pendulum: bool = False,
http_query_parameters: Optional[Sequence[Tuple[str, str]]] = None,
use_exact_imports: bool = False,
) -> None:
remote_text_cache: DefaultPutDict[str, str] = DefaultPutDict()
if isinstance(input_, str):
Expand Down Expand Up @@ -461,6 +462,7 @@ def get_header_and_first_line(csv_file: IO[str]) -> Dict[str, Any]:
custom_formatters_kwargs=custom_formatters_kwargs,
use_pendulum=use_pendulum,
http_query_parameters=http_query_parameters,
use_exact_imports=use_exact_imports,
**kwargs,
)

Expand Down
2 changes: 2 additions & 0 deletions datamodel_code_generator/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -312,6 +312,7 @@ def validate_root(cls, values: Any) -> Any:
custom_formatters_kwargs: Optional[TextIOBase] = None
use_pendulum: bool = False
http_query_parameters: Optional[Sequence[Tuple[str, str]]] = None
use_exact_imports: bool = False

def merge_args(self, args: Namespace) -> None:
set_args = {
Expand Down Expand Up @@ -508,6 +509,7 @@ def main(args: Optional[Sequence[str]] = None) -> Exit:
custom_formatters_kwargs=custom_formatters_kwargs,
use_pendulum=config.use_pendulum,
http_query_parameters=config.http_query_parameters,
use_exact_imports=config.use_exact_imports,
)
return Exit.OK
except InvalidClassNameError as e:
Expand Down
7 changes: 7 additions & 0 deletions datamodel_code_generator/arguments.py
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,13 @@ def start_section(self, heading: Optional[str]) -> None:
action='store_true',
default=False,
)
model_options.add_argument(
'--use-exact-imports',
help='import exact types instead of modules, for example: "from .foo import Bar" instead of '
'"from . import foo" with "foo.Bar"',
action='store_true',
default=False,
)

# ======================================================================================
# Typing options for generated models
Expand Down
3 changes: 2 additions & 1 deletion datamodel_code_generator/imports.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,11 +26,12 @@ class Imports(DefaultDict[Optional[str], Set[str]]):
def __str__(self) -> str:
return self.dump()

def __init__(self) -> None:
def __init__(self, use_exact: bool = False) -> None:
super().__init__(set)
self.alias: DefaultDict[Optional[str], Dict[str, str]] = defaultdict(dict)
self.counter: Dict[Tuple[Optional[str], str], int] = defaultdict(int)
self.reference_paths: Dict[str, Import] = {}
self.use_exact: bool = use_exact
alpoi-x marked this conversation as resolved.
Show resolved Hide resolved

def _set_alias(self, from_: Optional[str], imports: Set[str]) -> List[str]:
return [
Expand Down
18 changes: 16 additions & 2 deletions datamodel_code_generator/parser/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -238,6 +238,14 @@
return left, right


def exact_import(from_: str, import_: str, short_name: str) -> Tuple[str, str]:
if from_ == '.':
# Prevents "from . import foo" becoming "from ..foo import Foo"
# when our imported module has the same parent
return f'.{import_}', short_name
return f'{from_}.{import_}', short_name


@runtime_checkable
class Child(Protocol):
@property
Expand Down Expand Up @@ -392,6 +400,7 @@
custom_formatters_kwargs: Optional[Dict[str, Any]] = None,
use_pendulum: bool = False,
http_query_parameters: Optional[Sequence[Tuple[str, str]]] = None,
use_exact_imports: bool = False,
) -> None:
self.data_type_manager: DataTypeManager = data_type_manager_type(
python_version=target_python_version,
Expand All @@ -405,7 +414,8 @@
self.data_model_root_type: Type[DataModel] = data_model_root_type
self.data_model_field_type: Type[DataModelFieldBase] = data_model_field_type

self.imports: Imports = Imports()
self.imports: Imports = Imports(use_exact_imports)
self.use_exact_imports: bool = use_exact_imports
self._append_additional_imports(additional_imports=additional_imports)

self.base_class: Optional[str] = base_class
Expand Down Expand Up @@ -700,6 +710,10 @@
from_, import_ = full_path = relative(
model.module_name, data_type.full_name
)
if imports.use_exact:
from_, import_ = exact_import(

Check warning on line 714 in datamodel_code_generator/parser/base.py

View check run for this annotation

Codecov / codecov/patch

datamodel_code_generator/parser/base.py#L714

Added line #L714 was not covered by tests
from_, import_, data_type.reference.short_name
)
Comment on lines +713 to +716
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The repo prefer cover almost lines by the unittest.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hiya @koxudaxi, I'm not entirely sure how to test these lines 😓 I can't find any existing tests for this class method (Parser.__change_from_import) so I've got nothing to base it on.

Copy link
Contributor

@lord-haffi lord-haffi Jun 17, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@alpoi-x I think there are tests for executing the whole command. You could just take one of these and test the new option with multi-file-output (I don't remember where this test is but I think, I saw it few months ago). You could then assert that the imports look the way you want.

import_ = import_.replace('-', '_')

alias = scoped_model_resolver.add(full_path, import_).name
Expand Down Expand Up @@ -1250,7 +1264,7 @@
processed_models: List[Processed] = []

for module, models in module_models:
imports = module_to_import[module] = Imports()
imports = module_to_import[module] = Imports(self.use_exact_imports)
init = False
if module:
parent = (*module[:-1], '__init__.py')
Expand Down
2 changes: 2 additions & 0 deletions datamodel_code_generator/parser/graphql.py
Original file line number Diff line number Diff line change
Expand Up @@ -158,6 +158,7 @@ def __init__(
custom_formatters_kwargs: Optional[Dict[str, Any]] = None,
use_pendulum: bool = False,
http_query_parameters: Optional[Sequence[Tuple[str, str]]] = None,
use_exact_imports: bool = False,
) -> None:
super().__init__(
source=source,
Expand Down Expand Up @@ -225,6 +226,7 @@ def __init__(
custom_formatters_kwargs=custom_formatters_kwargs,
use_pendulum=use_pendulum,
http_query_parameters=http_query_parameters,
use_exact_imports=use_exact_imports,
)

self.data_model_scalar_type = data_model_scalar_type
Expand Down
2 changes: 2 additions & 0 deletions datamodel_code_generator/parser/jsonschema.py
Original file line number Diff line number Diff line change
Expand Up @@ -440,6 +440,7 @@ def __init__(
custom_formatters_kwargs: Optional[Dict[str, Any]] = None,
use_pendulum: bool = False,
http_query_parameters: Optional[Sequence[Tuple[str, str]]] = None,
use_exact_imports: bool = False,
) -> None:
super().__init__(
source=source,
Expand Down Expand Up @@ -507,6 +508,7 @@ def __init__(
custom_formatters_kwargs=custom_formatters_kwargs,
use_pendulum=use_pendulum,
http_query_parameters=http_query_parameters,
use_exact_imports=use_exact_imports,
)

self.remote_object_cache: DefaultPutDict[str, Dict[str, Any]] = DefaultPutDict()
Expand Down
2 changes: 2 additions & 0 deletions datamodel_code_generator/parser/openapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,7 @@ def __init__(
custom_formatters_kwargs: Optional[Dict[str, Any]] = None,
use_pendulum: bool = False,
http_query_parameters: Optional[Sequence[Tuple[str, str]]] = None,
use_exact_imports: bool = False,
):
super().__init__(
source=source,
Expand Down Expand Up @@ -289,6 +290,7 @@ def __init__(
custom_formatters_kwargs=custom_formatters_kwargs,
use_pendulum=use_pendulum,
http_query_parameters=http_query_parameters,
use_exact_imports=use_exact_imports,
)
self.open_api_scopes: List[OpenAPIScope] = openapi_scopes or [
OpenAPIScope.Schemas
Expand Down
20 changes: 19 additions & 1 deletion tests/parser/test_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,12 @@

from datamodel_code_generator.model import DataModel, DataModelFieldBase
from datamodel_code_generator.model.pydantic import BaseModel, DataModelField
from datamodel_code_generator.parser.base import Parser, relative, sort_data_models
from datamodel_code_generator.parser.base import (
Parser,
exact_import,
relative,
sort_data_models,
)
from datamodel_code_generator.reference import Reference, snake_to_upper_camel
from datamodel_code_generator.types import DataType

Expand Down Expand Up @@ -183,6 +188,19 @@ def test_relative(current_module: str, reference: str, val: Tuple[str, str]):
assert relative(current_module, reference) == val


@pytest.mark.parametrize(
'from_,import_,name,val',
[
('.', 'mod', 'Foo', ('.mod', 'Foo')),
('.a', 'mod', 'Foo', ('.a.mod', 'Foo')),
('..a', 'mod', 'Foo', ('..a.mod', 'Foo')),
('..a.b', 'mod', 'Foo', ('..a.b.mod', 'Foo')),
],
)
def test_exact_import(from_: str, import_: str, name: str, val: Tuple[str, str]):
assert exact_import(from_, import_, name) == val


@pytest.mark.parametrize(
'word,expected',
[
Expand Down
Loading