diff --git a/README.md b/README.md index 4edfa92..ac54f51 100644 --- a/README.md +++ b/README.md @@ -69,32 +69,43 @@ $ pip install usort To format one or more files or directories in-place: ```shell-session -$ usort format [ ...] +$ usort format [--config config/usort.toml] [ ...] ``` To generate a diff of changes without modifying files: ```shell-session -$ usort diff +$ usort diff [--config config/usort.toml] ``` To just validate that files are formatted correctly, like during CI: ```shell-session -$ usort check +$ usort check [--config config/usort.toml] +``` + +### Explicit configuration files + +All CLI commands accept `--config` to point at an alternate TOML configuration file. +If omitted µsort falls back to the discovery of the nearest `pyproject.toml`. + +```shell-session +$ usort format --config config/usort.toml ``` ### pre-commit µsort provides a [pre-commit](https://pre-commit.com/) hook. To enforce sorted imports before every commit, add the following to your `.pre-commit-config.yaml` -file: +file. If you keep µsort settings in a custom TOML file, pass the same +`--config` flag via `args`: ```yaml - repo: https://github.com/facebook/usort rev: v1.0.7 hooks: - id: usort + args: ["--config", "config/usort.toml"] ``` ## License diff --git a/docs/guide.rst b/docs/guide.rst index 636ecee..c867a39 100644 --- a/docs/guide.rst +++ b/docs/guide.rst @@ -12,19 +12,19 @@ To format one or more files or directories in-place: .. code-block:: shell-session - $ usort format [ ...] + $ usort format [--config config/usort.toml] [ ...] To generate a diff of changes without modifying files: .. code-block:: shell-session - $ usort diff [ ...] + $ usort diff [--config config/usort.toml] [ ...] µsort can also be used to validate formatting as part of CI: .. code-block:: shell-session - $ usort check [ ...] + $ usort check [--config config/usort.toml] [ ...] Sorting @@ -462,6 +462,17 @@ When sorting each file, µsort will look for the "nearest" :file:`pyproject.toml to the file being sorted, looking upwards until the project root is found, or until the root of the filesystem is reached. +Explicit config files +%%%%%%%%%%%%%%%%%%%%% + +You can specify a configuration file using the ``--config`` CLI option to point +commands like ``format``, ``check``, ``diff`` and ``list-imports`` at a specific +TOML file: + +.. code-block:: shell-session + + $ usort format --config config/usort.toml src/ + ``[tool.usort]`` %%%%%%%%%%%%%%%% diff --git a/usort/api.py b/usort/api.py index 329395a..cc356a3 100644 --- a/usort/api.py +++ b/usort/api.py @@ -91,7 +91,9 @@ def usort_string(data: str, config: Config, path: Optional[Path] = None) -> str: return result.output.decode() -def usort_file(path: Path, *, write: bool = False) -> Result: +def usort_file( + path: Path, *, write: bool = False, config: Optional[Config] = None +) -> Result: """ Format a single file and return a Result object. @@ -99,7 +101,13 @@ def usort_file(path: Path, *, write: bool = False) -> Result: """ try: - config = Config.find(path.parent) + if config is None: + config = Config.find(path.parent) + + # Apply first-party detection when enabled + if config.first_party_detection: + config = config.with_first_party(Path.cwd() / path.parent) + data = path.read_bytes() result = usort(data, config, path) @@ -117,7 +125,9 @@ def usort_file(path: Path, *, write: bool = False) -> Result: def usort_path( - paths: Union[Path, Iterable[Path]], write: bool = False + paths: Union[Path, Iterable[Path]], + write: bool = False, + config: Optional[Config] = None, ) -> Iterable[Result]: """ For a given path, format it, or any python files in it, and yield :class:`Result` s. @@ -134,12 +144,19 @@ def usort_path( with timed("total"): materialized_paths: Set[Path] = set() - for path in source_paths: - with timed(f"walking {path}"): - config = Config.find(path) - materialized_paths.update(walk(path, excludes=config.excludes)) + if config is None: + # discover config per-root + for path in source_paths: + with timed(f"walking {path}"): + discovered = Config.find(path) + materialized_paths.update(walk(path, excludes=discovered.excludes)) + else: + # use global config + for path in source_paths: + with timed(f"walking {path}"): + materialized_paths.update(walk(path, excludes=config.excludes)) - fn = partial(usort_file, write=write) + fn = partial(usort_file, write=write, config=config) if len(materialized_paths) == 1: # shave off multiprocessing overhead results = [fn(materialized_paths.pop())] @@ -148,7 +165,7 @@ def usort_path( return results -def usort_stdin() -> bool: +def usort_stdin(config: Optional[Config] = None) -> bool: """ Read file contents from stdin, format it, and write the resulting file to stdout @@ -161,7 +178,10 @@ def usort_stdin() -> bool: print("Warning: stdin is a tty", file=sys.stderr) try: - config = Config.find() + if config is None: + config = Config.find() + if config.first_party_detection: + config = config.with_first_party(Path.cwd()) data = sys.stdin.read() result = usort_string(data, config, Path("")) sys.stdout.write(result) diff --git a/usort/cli.py b/usort/cli.py index d7e23bb..3e5bae9 100644 --- a/usort/cli.py +++ b/usort/cli.py @@ -7,12 +7,12 @@ import sys from functools import wraps from pathlib import Path -from typing import Any, Callable, List, Sequence +from typing import Any, Callable, List, Optional, Sequence import click from moreorless.click import echo_color_unified_diff - from usort.translate import render_node + from . import __version__ from .api import usort_path, usort_stdin from .config import Config @@ -43,6 +43,15 @@ def wrapper(*args: Any, **kwargs: Any) -> None: return wrapper +def _load_config(config_path: Optional[Path]) -> Optional[Config]: + if config_path is None: + return None + try: + return Config.load(config_path) + except Exception as e: + raise click.ClickException(str(e)) + + @click.group() @click.pass_context @click.version_option(__version__, "--version", "-V") @@ -70,17 +79,32 @@ def main(ctx: click.Context, benchmark: bool, debug: bool, native: bool) -> None @main.command() @click.option("--multiples", is_flag=True, help="Only show files with multiple blocks") @click.option("--debug", is_flag=True, help="Show internal information") +@click.option( + "--config", + "config_path", + type=click.Path(exists=True, dir_okay=False, path_type=Path), + help="Explicit TOML config file path", +) @click.argument("paths", nargs=-1, type=click.Path(dir_okay=False, path_type=Path)) @usort_command -def list_imports(multiples: bool, debug: bool, paths: List[Path]) -> int: +def list_imports( + multiples: bool, debug: bool, config_path: Optional[Path], paths: List[Path] +) -> int: """ Troubleshoot sorting behavior and show import blocks """ # This is used to debug the sort keys on the various lines, and understand # where the barriers are that produce different blocks. + explicit_config = _load_config(config_path) + for path in paths: - config = Config.find(path) + config = explicit_config or Config.find(path) + + # Apply first-party detection when enabled + if config.first_party_detection: + config = config.with_first_party(Path.cwd() / path) + mod = try_parse(path) try: sorter = ImportSorter(module=mod, path=path, config=config) @@ -114,9 +138,15 @@ def list_imports(multiples: bool, debug: bool, paths: List[Path]) -> int: @main.command() +@click.option( + "--config", + "config_path", + type=click.Path(exists=True, dir_okay=False, path_type=Path), + help="Explicit TOML config file path", +) @click.argument("paths", nargs=-1, type=click.Path(path_type=Path)) @usort_command -def check(paths: List[Path]) -> int: +def check(paths: List[Path], config_path: Optional[Path]) -> int: """ Check imports for one or more path """ @@ -124,7 +154,8 @@ def check(paths: List[Path]) -> int: raise click.ClickException("Provide some filenames") return_code = 0 - for result in usort_path(paths, write=False): + config = _load_config(config_path) + for result in usort_path(paths, write=False, config=config): if result.error: click.echo(f"Error sorting {result.path}: {result.error}") return_code |= 1 @@ -142,10 +173,16 @@ def check(paths: List[Path]) -> int: @main.command() +@click.option( + "--config", + "config_path", + type=click.Path(exists=True, dir_okay=False, path_type=Path), + help="Explicit TOML config file path", +) @click.pass_context @click.argument("paths", nargs=-1, type=click.Path(path_type=Path)) @usort_command -def diff(ctx: click.Context, paths: List[Path]) -> int: +def diff(ctx: click.Context, config_path: Optional[Path], paths: List[Path]) -> int: """ Output diff of changes for one or more path """ @@ -153,7 +190,8 @@ def diff(ctx: click.Context, paths: List[Path]) -> int: raise click.ClickException("Provide some filenames") return_code = 0 - for result in usort_path(paths, write=False): + config = _load_config(config_path) + for result in usort_path(paths, write=False, config=config): if result.error: click.echo(f"Error sorting {result.path}: {result.error}") if ctx.obj.debug: @@ -178,9 +216,15 @@ def diff(ctx: click.Context, paths: List[Path]) -> int: @main.command() +@click.option( + "--config", + "config_path", + type=click.Path(exists=True, dir_okay=False, path_type=Path), + help="Explicit TOML config file path", +) @click.argument("paths", nargs=-1, type=click.Path(allow_dash=True, path_type=Path)) @usort_command -def format(paths: List[Path]) -> int: +def format(paths: List[Path], config_path: Optional[Path]) -> int: """ Format one or more paths @@ -191,12 +235,13 @@ def format(paths: List[Path]) -> int: if not paths: raise click.ClickException("Provide some filenames") + config = _load_config(config_path) if paths[0] == Path("-"): - success = usort_stdin() + success = usort_stdin(config=config) return 0 if success else 1 return_code = 0 - for result in usort_path(paths, write=True): + for result in usort_path(paths, write=True, config=config): if result.error: click.echo(f"Error sorting {result.path}: {result.error}") return_code |= 1 diff --git a/usort/config.py b/usort/config.py index 366ccff..8184f51 100644 --- a/usort/config.py +++ b/usort/config.py @@ -8,6 +8,7 @@ from dataclasses import dataclass, field from pathlib import Path from typing import Dict, List, NewType, Optional, Pattern, Sequence, Set +from warnings import warn if sys.version_info < (3, 11): import tomli as tomllib @@ -56,7 +57,7 @@ class Config: side_effect_modules: List[str] = field(default_factory=list) side_effect_re: Pattern[str] = field(default=re.compile("")) - # Whether to perform the first-party heuristic during find() + # Whether to perform the first-party heuristic first_party_detection: bool = True # Whether to follow black-style for magic trailing commas @@ -85,9 +86,10 @@ def __post_init__(self) -> None: ) @classmethod - def find( - cls, filename: Optional[Path] = None, with_first_party: bool = True - ) -> "Config": + def find(cls, filename: Optional[Path] = None) -> "Config": + """ + Find and load configuration by walking up from the given path. + """ rv = cls() # TODO This logic should be split out to a separate project, as it's @@ -114,15 +116,6 @@ def find( p = p.parent - # Either param or config can force off. - if with_first_party and rv.first_party_detection: - if filename is None: - # directory - rv = rv.with_first_party(Path.cwd()) - else: - # filename, ideally - rv = rv.with_first_party(Path.cwd() / filename) - return rv def with_first_party(self, filename: Path) -> "Config": @@ -212,6 +205,26 @@ def update_from_config(self, toml_path: Path) -> None: # make sure generated regexes get updated self.__post_init__() + @classmethod + def load(cls, path: Path) -> "Config": + """Load a config from an explicit TOML file path. + + Warns if the file does not contain a [tool.usort] or [tool.black] table. + """ + if not path.exists(): + raise FileNotFoundError(f"Config file not found: {path}") + + text = path.read_text() + conf = tomllib.loads(text) + tool = conf.get("tool", {}) + if "tool" not in conf or ("usort" not in tool and "black" not in tool): + warn(f"Config file {path} is missing [tool.usort] or [tool.black] table") + + rv = cls() + # Re-parse via update_from_config to leverage existing logic + rv.update_from_config(path) + return rv + def category(self, dotted_import: str) -> Category: """ Given a piece of an import string, return its category for this config. diff --git a/usort/tests/__init__.py b/usort/tests/__init__.py index 67dfb7d..b678b58 100644 --- a/usort/tests/__init__.py +++ b/usort/tests/__init__.py @@ -5,7 +5,12 @@ from .cli import CliTest from .config import ConfigTest -from .functional import BasicOrderingTest, UsortStringFunctionalTest +from .functional import ( + BasicOrderingTest, + UsortFileFirstPartyTest, + UsortStdinFirstPartyTest, + UsortStringFunctionalTest, +) from .sorting import SplitTest from .stdlibs import StdlibsTest from .translate import IsSortableTest, SortableImportTest @@ -16,6 +21,8 @@ "CliTest", "ConfigTest", "BasicOrderingTest", + "UsortFileFirstPartyTest", + "UsortStdinFirstPartyTest", "UsortStringFunctionalTest", "IsSortableTest", "SortableImportTest", diff --git a/usort/tests/cli.py b/usort/tests/cli.py index c99dd1f..a4421e3 100644 --- a/usort/tests/cli.py +++ b/usort/tests/cli.py @@ -211,6 +211,83 @@ def test_list_imports(self) -> None: ) self.assertEqual(result.exit_code, 0) + def test_list_imports_auto_discovered_config(self) -> None: + # Test that list-imports respects first_party_detection with auto-discovered config + import tempfile + + with tempfile.TemporaryDirectory() as d: + (Path(d) / "a" / "b" / "c").mkdir(parents=True) + (Path(d) / "a" / "b" / "__init__.py").write_text("") + (Path(d) / "a" / "b" / "c" / "__init__.py").write_text( + "import sys\nimport b\nimport other\n" + ) + + runner = CliRunner() + with chdir(str(Path(d) / "a" / "b")): + result = runner.invoke(main, ["list-imports", "c/__init__.py"]) + + # b should be detected as first-party and sorted into its own block + self.assertIn("import b", result.output) + self.assertIn("import other", result.output) + self.assertIn("import sys", result.output) + self.assertEqual(result.exit_code, 0) + + # Now test with first_party_detection disabled in auto-discovered config + (Path(d) / "pyproject.toml").write_text( + "[tool.usort]\nfirst_party_detection = false\n" + ) + + with chdir(str(Path(d) / "a" / "b")): + result = runner.invoke(main, ["list-imports", "c/__init__.py"]) + + # Without first-party detection, b and other should be in the same block + self.assertIn("import b", result.output) + self.assertIn("import other", result.output) + self.assertEqual(result.exit_code, 0) + + def test_list_imports_explicit_config(self) -> None: + # Test that list-imports respects first_party_detection with explicit config + import tempfile + + with tempfile.TemporaryDirectory() as d: + (Path(d) / "a" / "b" / "c").mkdir(parents=True) + (Path(d) / "a" / "b" / "__init__.py").write_text("") + (Path(d) / "a" / "b" / "c" / "__init__.py").write_text( + "import sys\nimport b\nimport other\n" + ) + + # Config with first_party_detection enabled (default) + config_enabled = Path(d) / "enabled.toml" + config_enabled.write_text("[tool.usort]\n") + + runner = CliRunner() + with chdir(str(Path(d) / "a" / "b")): + result = runner.invoke( + main, ["list-imports", "--config", str(config_enabled), "c/__init__.py"] + ) + + # b should be detected as first-party + self.assertIn("import b", result.output) + self.assertIn("import other", result.output) + self.assertIn("import sys", result.output) + self.assertEqual(result.exit_code, 0) + + # Config with first_party_detection disabled + config_disabled = Path(d) / "disabled.toml" + config_disabled.write_text( + "[tool.usort]\nfirst_party_detection = false\n" + ) + + with chdir(str(Path(d) / "a" / "b")): + result = runner.invoke( + main, ["list-imports", "--config", str(config_disabled), "c/__init__.py"] + ) + + # Without first-party detection, b and other should be in the same block + self.assertIn("import b", result.output) + self.assertIn("import other", result.output) + self.assertEqual(result.exit_code, 0) + def test_format_no_change(self) -> None: with sample_contents("import sys\n") as dtmp: runner = CliRunner() @@ -372,3 +449,105 @@ def test_format_latin_1(self) -> None: ), # git on windows again (Path(dtmp) / "sample.py").read_bytes(), ) + + def test_diff_explicit_config(self) -> None: + # Test that diff command respects explicit config + with volatile.dir() as dtmp: + root = Path(dtmp) + (root / "pyproject.toml").write_text( + """ +[tool.usort] +merge_imports = true +""".strip() + ) + (root / "altconfig.toml").write_text( + """ +[tool.usort] +merge_imports = false +""".strip() + ) + (root / "sample.py").write_text("from foo import b\nfrom foo import a\n") + runner = CliRunner() + with chdir(dtmp): + # With explicit config (merge disabled), should show diff for separate imports + result = runner.invoke( + main, ["diff", "--config", "altconfig.toml", "sample.py"] + ) + self.assertEqual(result.exit_code, 0) # diff returns 0 even with changes + # Should show changes (sorting a before b, not merging) + self.assertIn("from foo import a", result.output) + self.assertIn("from foo import b", result.output) + # Verify it's actually using the explicit config by checking output format + self.assertNotIn("a, b", result.output) # Not merged + + def test_explicit_config_override_merge_imports_disabled(self) -> None: + # pyproject enables merging, explicit config disables it + with volatile.dir() as dtmp: + root = Path(dtmp) + (root / "pyproject.toml").write_text( + """ +[tool.usort] +merge_imports = true +""".strip() + ) + (root / "altconfig.toml").write_text( + """ +[tool.usort] +merge_imports = false +""".strip() + ) + (root / "sample.py").write_text("from foo import b\nfrom foo import a\n") + runner = CliRunner() + with chdir(dtmp): + # Without override: merged into single statement + result_default = runner.invoke(main, ["format", "."]) + self.assertEqual(result_default.exit_code, 0) + self.assertEqual( + (root / "sample.py").read_text(), + "from foo import a, b\n", + ) + # Reset contents + (root / "sample.py").write_text("from foo import b\nfrom foo import a\n") + with chdir(dtmp): + result_override = runner.invoke( + main, ["format", "--config", "altconfig.toml", "."] + ) + self.assertEqual(result_override.exit_code, 0) + # With override: statements remain separate (but sorted) + self.assertEqual( + (root / "sample.py").read_text(), + "from foo import a\nfrom foo import b\n", + ) + + def test_explicit_config_missing_table_warning(self) -> None: + # Config files without tool.usort or tool.black should warn but still work + with volatile.dir() as dtmp: + root = Path(dtmp) + (root / "pyproject.toml").write_text( + """ +[tool.usort] +merge_imports = true +""".strip() + ) + (root / "altconfig.toml").write_text("# no tool.usort table\n") + (root / "sample.py").write_text("import sys\n") + runner = CliRunner() + with chdir(dtmp): + result = runner.invoke( + main, ["check", "--config", "altconfig.toml", "."] + ) + # Should succeed (exit code 0) with a warning + self.assertEqual(result.exit_code, 0) + + def test_explicit_config_nonexistent_file(self) -> None: + # Non-existent config file should error + with volatile.dir() as dtmp: + root = Path(dtmp) + (root / "sample.py").write_text("import sys\n") + runner = CliRunner() + with chdir(dtmp): + result = runner.invoke( + main, ["check", "--config", "nonexistent.toml", "."] + ) + self.assertNotEqual(result.exit_code, 0) + self.assertIn("does not exist", result.output.lower()) diff --git a/usort/tests/config.py b/usort/tests/config.py index 9bbf241..b589e34 100644 --- a/usort/tests/config.py +++ b/usort/tests/config.py @@ -8,6 +8,7 @@ from pathlib import Path from usort.config import CAT_FIRST_PARTY, CAT_FUTURE, CAT_THIRD_PARTY, Config + from .cli import chdir @@ -60,44 +61,6 @@ def test_third_party(self) -> None: new_conf = Config.find(Path(d)) self.assertEqual(conf, new_conf) - def test_first_party_root_finding(self) -> None: - with tempfile.TemporaryDirectory() as d: - (Path(d) / "a" / "b" / "c").mkdir(parents=True) - (Path(d) / "a" / "b" / "__init__.py").write_text("") - (Path(d) / "a" / "b" / "c" / "__init__.py").write_text("from b import zzz") - - f = Path(d) / "a" / "b" / "c" / "__init__.py" - - conf = Config.find(f) - self.assertEqual(CAT_FIRST_PARTY, conf.known["b"]) - conf = Config.find(f.parent) # c - self.assertEqual(CAT_FIRST_PARTY, conf.known["b"]) - conf = Config.find(f.parent.parent) # b - self.assertEqual(CAT_FIRST_PARTY, conf.known["b"]) - conf = Config.find(Path("/")) - self.assertNotIn("b", conf.known) - - def test_first_party_root_finding_disable(self) -> None: - with tempfile.TemporaryDirectory() as d: - (Path(d) / "a" / "b" / "c").mkdir(parents=True) - (Path(d) / "a" / "b" / "__init__.py").write_text("") - (Path(d) / "a" / "b" / "c" / "__init__.py").write_text("from b import zzz") - - f = Path(d) / "a" / "b" / "c" / "__init__.py" - - conf = Config.find(f) - self.assertEqual(CAT_FIRST_PARTY, conf.known["b"]) - - (Path(d) / "pyproject.toml").write_text( - """\ -[tool.usort] -first_party_detection = false -""" - ) - - conf = Config.find(f) - self.assertNotIn("b", conf.known) - def test_new_category_names(self) -> None: with tempfile.TemporaryDirectory() as d: (Path(d) / "pyproject.toml").write_text( @@ -277,19 +240,84 @@ def test_config_parent_walk(self) -> None: (d_path / "b" / "c" / "d").symlink_to(d_path / "a") # This is behavior that's worked for ages, if given an absolute path - conf = Config.find(d_path / "b" / "c" / "d" / "__init__.py") + conf = Config.find( + d_path / "b" / "c" / "d" / "__init__.py" + ).with_first_party(Path.cwd() / d_path / "b" / "c" / "d" / "__init__.py") self.assertEqual(CAT_FIRST_PARTY, conf.known["x"]) self.assertEqual(CAT_FIRST_PARTY, conf.known["d"]) # This is also something that's worked for ages, a relative path with cwd # above the pyproject.toml with chdir(d): - conf = Config.find(Path("b") / "c" / "d" / "__init__.py") + conf = Config.find( + Path("b") / "c" / "d" / "__init__.py" + ).with_first_party(Path.cwd() / Path("b") / "c" / "d" / "__init__.py") self.assertEqual(CAT_FIRST_PARTY, conf.known["x"]) self.assertEqual(CAT_FIRST_PARTY, conf.known["d"]) # This is an instance of issue 43 with chdir((d_path / "b" / "c").as_posix()): - conf = Config.find(Path("d") / "__init__.py") + conf = Config.find(Path("d") / "__init__.py").with_first_party( + Path.cwd() / Path("d") / "__init__.py" + ) self.assertEqual(CAT_FIRST_PARTY, conf.known["x"]) self.assertEqual(CAT_FIRST_PARTY, conf.known["d"]) + + def test_load_explicit_config(self) -> None: + with tempfile.TemporaryDirectory() as d: + d_path = Path(d) + config_file = d_path / "custom.toml" + + with self.subTest("with tool.usort"): + config_file.write_text( + """ +[tool.usort] +merge_imports = false +known_third_party = ["pandas"] +""" + ) + conf = Config.load(config_file) + self.assertFalse(conf.merge_imports) + self.assertEqual(CAT_THIRD_PARTY, conf.known["pandas"]) + + with self.subTest("with tool.black only"): + config_file.write_text( + """ +[tool.black] +line-length = 100 +""" + ) + # Config should still work with black settings + conf = Config.load(config_file) + self.assertEqual(100, conf.line_length) + + with self.subTest("with both tool.usort and tool.black"): + config_file.write_text( + """ +[tool.usort] +merge_imports = false + +[tool.black] +line-length = 120 +""" + ) + conf = Config.load(config_file) + self.assertFalse(conf.merge_imports) + self.assertEqual(120, conf.line_length) + + with self.subTest("missing file"): + missing = d_path / "nonexistent.toml" + with self.assertRaises(FileNotFoundError): + Config.load(missing) + + with self.subTest("empty config file"): + config_file.write_text("# empty\n") + import warnings + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + conf = Config.load(config_file) + self.assertEqual(1, len(w)) + self.assertIn("missing", str(w[0].message).lower()) + # Should still create a valid config with defaults + self.assertTrue(conf.merge_imports) # default value diff --git a/usort/tests/functional.py b/usort/tests/functional.py index 5fff654..be72ec1 100644 --- a/usort/tests/functional.py +++ b/usort/tests/functional.py @@ -3,14 +3,17 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. +import os import tempfile import unittest from dataclasses import replace +from io import StringIO from pathlib import Path from textwrap import dedent from typing import Optional +from unittest.mock import patch -from ..api import usort, usort_path +from ..api import usort, usort_file, usort_path, usort_stdin from ..config import CAT_FIRST_PARTY, Config from ..translate import import_from_node from ..util import parse_import @@ -1323,5 +1326,158 @@ def test_collapse_blank_lines_in_category_false(self) -> None: ) +class UsortFileFirstPartyTest(unittest.TestCase): + """Test that usort_file respects first_party_detection with auto-discovered and explicit configs""" + + def test_auto_discovered_config(self) -> None: + with tempfile.TemporaryDirectory() as d: + (Path(d) / "a" / "b" / "c").mkdir(parents=True) + (Path(d) / "a" / "b" / "__init__.py").write_text("") + (Path(d) / "a" / "b" / "c" / "__init__.py").write_text( + "import sys\nimport b\nimport other\n" + ) + + # First verify default behavior (detection enabled) + result = usort_file(Path(d) / "a" / "b" / "c" / "__init__.py") + self.assertIsNone(result.error) + expected = "import sys\n\nimport other\n\nimport b\n" + self.assertEqual(result.output.decode("utf-8"), expected) + + # Now disable first-party detection in config + (Path(d) / "pyproject.toml").write_text( + """\ +[tool.usort] +first_party_detection = false +""" + ) + + # usort_file should NOT detect 'b' as first-party + result = usort_file(Path(d) / "a" / "b" / "c" / "__init__.py") + self.assertIsNone(result.error) + # b should be treated as third-party + expected = "import sys\n\nimport b\nimport other\n" + self.assertEqual(result.output.decode("utf-8"), expected) + + def test_explicit_config(self) -> None: + with tempfile.TemporaryDirectory() as d: + (Path(d) / "a" / "b" / "c").mkdir(parents=True) + (Path(d) / "a" / "b" / "__init__.py").write_text("") + (Path(d) / "a" / "b" / "c" / "__init__.py").write_text( + "import sys\nimport b\nimport other\n" + ) + + # Config with first_party_detection enabled (default) + config_enabled = Path(d) / "enabled.toml" + config_enabled.write_text("[tool.usort]\n") + + result = usort_file( + Path(d) / "a" / "b" / "c" / "__init__.py", + config=Config.load(config_enabled), + ) + self.assertIsNone(result.error) + # b should be in its own section (first-party) + expected = "import sys\n\nimport other\n\nimport b\n" + self.assertEqual(result.output.decode("utf-8"), expected) + + # Config with first_party_detection disabled + config_disabled = Path(d) / "disabled.toml" + config_disabled.write_text( + "[tool.usort]\nfirst_party_detection = false\n" + ) + + result = usort_file( + Path(d) / "a" / "b" / "c" / "__init__.py", + config=Config.load(config_disabled), + ) + self.assertIsNone(result.error) + # b should NOT be detected as first-party, treated as third-party + expected = "import sys\n\nimport b\nimport other\n" + self.assertEqual(result.output.decode("utf-8"), expected) + + +class UsortStdinFirstPartyTest(unittest.TestCase): + """Test that usort_stdin respects first_party_detection with auto-discovered and explicit configs""" + + def test_auto_discovered_config(self) -> None: + # Test that usort_stdin respects first_party_detection with auto-discovered config + with tempfile.TemporaryDirectory() as d: + (Path(d) / "a" / "b" / "c").mkdir(parents=True) + (Path(d) / "a" / "b" / "__init__.py").write_text("") + + # Simulate stdin with imports + input_text = "import sys\nimport b\nimport other\n" + + orig_cwd = os.getcwd() + try: + os.chdir(Path(d) / "a" / "b") + + # Default behavior: auto-detect first-party + with patch("sys.stdin", StringIO(input_text)): + with patch("sys.stdout", new_callable=StringIO) as mock_stdout: + result = usort_stdin() + + self.assertTrue(result) + # b should be detected as first-party + expected = "import sys\n\nimport other\n\nimport b\n" + self.assertEqual(mock_stdout.getvalue(), expected) + + # Create config with first_party_detection disabled + (Path(d) / "pyproject.toml").write_text( + "[tool.usort]\nfirst_party_detection = false\n" + ) + + with patch("sys.stdin", StringIO(input_text)): + with patch("sys.stdout", new_callable=StringIO) as mock_stdout: + result = usort_stdin() + + self.assertTrue(result) + # b should NOT be detected as first-party + expected = "import sys\n\nimport b\nimport other\n" + self.assertEqual(mock_stdout.getvalue(), expected) + finally: + os.chdir(orig_cwd) + + def test_explicit_config(self) -> None: + with tempfile.TemporaryDirectory() as d: + (Path(d) / "a" / "b" / "c").mkdir(parents=True) + (Path(d) / "a" / "b" / "__init__.py").write_text("") + + # Config with first_party_detection enabled (default) + config_enabled = Path(d) / "enabled.toml" + config_enabled.write_text("[tool.usort]\n") + + # Simulate stdin with imports + input_text = "import sys\nimport b\nimport other\n" + + orig_cwd = os.getcwd() + try: + os.chdir(Path(d) / "a" / "b") + with patch("sys.stdin", StringIO(input_text)): + with patch("sys.stdout", new_callable=StringIO) as mock_stdout: + result = usort_stdin(config=Config.load(config_enabled)) + + self.assertTrue(result) + # b should be detected as first-party + expected = "import sys\n\nimport other\n\nimport b\n" + self.assertEqual(mock_stdout.getvalue(), expected) + + # Config with first_party_detection disabled + config_disabled = Path(d) / "disabled.toml" + config_disabled.write_text( + "[tool.usort]\nfirst_party_detection = false\n" + ) + + with patch("sys.stdin", StringIO(input_text)): + with patch("sys.stdout", new_callable=StringIO) as mock_stdout: + result = usort_stdin(config=Config.load(config_disabled)) + + self.assertTrue(result) + # b should NOT be detected as first-party + expected = "import sys\n\nimport b\nimport other\n" + self.assertEqual(mock_stdout.getvalue(), expected) + finally: + os.chdir(orig_cwd) + + if __name__ == "__main__": unittest.main()