Merge pull request #18 from souliane/refactor-get-poetry-deps

poetry_to_pre_commit/common.py

@@ -23,10 +23,18 @@ def pre_commit_config_roundtrip(
         yaml.dump(config, path)
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(frozen=True, order=True)
 class PoetryPackage:
     name: str
-    version: str
+    version: str = "*"
+    extras: frozenset[str] = dataclasses.field(default_factory=frozenset)
+
+    def __str__(self) -> str:
+        if self.extras:
+            extras = ",".join(sorted(self.extras))
+            return f"{self.name}[{extras}]=={self.version}"
+        else:
+            return f"{self.name}=={self.version}"
 
 
 def get_poetry_packages(cwd: pathlib.Path | None = None) -> Iterable[PoetryPackage]:
@@ -35,6 +43,5 @@ def get_poetry_packages(cwd: pathlib.Path | None = None) -> Iterable[PoetryPacka
 
     for package in repository.packages:
         yield PoetryPackage(
-            name=package.name,
-            version=package.version.text,
+            name=package.name, version=package.version.text, extras=package.features
         )

poetry_to_pre_commit/sync_hooks_additional_dependencies.py

@@ -55,7 +55,9 @@ def get_sync_hooks_additional_dependencies_parser() -> argparse.ArgumentParser:
     return parser
 
 
-def get_poetry_deps(*, cwd: pathlib.Path | None = None, group: str) -> Iterable[str]:
+def get_poetry_deps(
+    *, cwd: pathlib.Path | None = None, group: str
+) -> Iterable[common.PoetryPackage]:
     package_by_name = {p.name: p for p in common.get_poetry_packages(cwd=cwd)}
     main_package = factory.Factory().create_poetry(cwd=cwd).package
     try:
@@ -71,28 +73,25 @@ def get_poetry_deps(*, cwd: pathlib.Path | None = None, group: str) -> Iterable[
                 f"Package not found in poetry.lock: {dep.name}. "
                 "Is your poetry.lock up-to-date?"
             ) from e
-        yield f"{dep.complete_name}=={package.version}"
+        yield common.PoetryPackage(
+            name=dep.name, version=package.version, extras=dep.extras
+        )
 
 
 def update_or_remove_additional_deps(
-    poetry_deps: set[str], hook_additional_deps: list[str]
-) -> set[str]:
+    poetry_deps: set[common.PoetryPackage], hook_additional_deps: list[str]
+) -> set[common.PoetryPackage]:
     # Additional packages that are already in pre-commit configuration could be listed with
     # any format that is accepted by pip - use `Requirement` to parse them properly.
     current_deps = [Requirement(dep).name for dep in hook_additional_deps]
 
-    return {
-        package
-        for package in poetry_deps
-        # package is yielded by `get_poetry_deps` above, and we are pretty sure that this won't raise `IndexError`
-        if package.split("==")[0].split("[")[0] in current_deps
-    }
+    return {package for package in poetry_deps if package.name in current_deps}
 
 
 def _sync_hooks_additional_dependencies(
     *,
     config: dict[str, Any],
-    deps_by_group: dict[str, set[str]],
+    deps_by_group: dict[str, set[common.PoetryPackage]],
     bind: dict[str, set[str]],
     no_new_deps: bool = False,
 ) -> None:
@@ -112,19 +111,22 @@ def _sync_hooks_additional_dependencies(
                 groups = bind[hook_id]
             except KeyError:
                 continue
-            deps: set[str] = set()
+            deps: set[common.PoetryPackage] = set()
 
             for group in groups:
                 deps.update(deps_by_group.get(group, set()))
 
-            hook["additional_dependencies"] = sorted(
+            packages = (
                 update_or_remove_additional_deps(
                     poetry_deps=deps,
                     hook_additional_deps=hook["additional_dependencies"],
                 )
                 if no_new_deps
                 else deps
             )
+            hook["additional_dependencies"] = sorted(
+                str(package) for package in packages
+            )
 
 
 def sync_hooks_additional_dependencies(
@@ -137,7 +139,7 @@ def sync_hooks_additional_dependencies(
     args = parser.parse_args(argv)
 
     bind = combine_bind_values(args.bind)
-    deps_by_group: dict[str, set[str]] = {}
+    deps_by_group: dict[str, set[common.PoetryPackage]] = {}
 
     for groups in bind.values():
         for group in groups:

tests/test_sync_hooks_additional_dependencies.py

@@ -6,6 +6,7 @@
 import ruamel.yaml
 
 from poetry_to_pre_commit import sync_hooks_additional_dependencies
+from poetry_to_pre_commit.common import PoetryPackage
 
 
 @pytest.mark.parametrize(
@@ -33,9 +34,7 @@ def test_combine_bind_values() -> None:
 
 
 def test_get_sync_hooks_additional_dependencies_parser() -> None:
-    parser = (
-        sync_hooks_additional_dependencies.get_sync_hooks_additional_dependencies_parser()
-    )
+    parser = sync_hooks_additional_dependencies.get_sync_hooks_additional_dependencies_parser()
     assert parser.parse_args(["--bind", "foo=bar,baz", "--bind", "foo=qux"]).bind == [
         ("foo", {"bar", "baz"}),
         ("foo", {"qux"}),
@@ -51,9 +50,9 @@ def test_get_poetry_deps(poetry_cwd: Path) -> None:
     )
 
     assert sorted(results) == [
-        "flake8==7.0.0",
-        "pylint==3.1.0",
-        "pyright==1.1.355",
+        PoetryPackage("flake8", "7.0.0"),
+        PoetryPackage("pylint", "3.1.0"),
+        PoetryPackage("pyright", "1.1.355"),
     ]
 
 
@@ -130,16 +129,29 @@ def test_sync_hooks_additional_dependencies(tmp_path: Path, poetry_cwd: Path) ->
 @pytest.mark.parametrize(
     ("poetry_deps", "additional_deps", "expected_additional_deps"),
     [
-        (["a==1", "b"], ["a"], ["a==1"]),
-        (["a==1", "b"], ["a[x]==2"], ["a==1"]),
-        (["a[x]==1", "b"], ["a"], ["a[x]==1"]),
-        (["a[x]==1", "b"], ["a[x]"], ["a[x]==1"]),
-        (["a==1", "b"], ["a == 2"], ["a==1"]),
-        (["a==1", "b"], ["a<=2"], ["a==1"]),
-        (["a==1", "b"], ["a>=1"], ["a==1"]),
+        ({PoetryPackage("a", "1"), PoetryPackage("b")}, ["a"], ["a==1"]),
+        ({PoetryPackage("a", "1"), PoetryPackage("b")}, ["a[x]==2"], ["a==1"]),
+        (
+            {PoetryPackage("a", "1", frozenset({"x"})), PoetryPackage("b")},
+            ["a"],
+            ["a[x]==1"],
+        ),
+        (
+            {PoetryPackage("a", "1", frozenset({"x"})), PoetryPackage("b")},
+            ["a[x]"],
+            ["a[x]==1"],
+        ),
+        (
+            {PoetryPackage("a", "1", frozenset({"x", "y"})), PoetryPackage("b")},
+            ["a[x]"],
+            ["a[x,y]==1"],
+        ),
+        ({PoetryPackage("a", "1"), PoetryPackage("b")}, ["a == 2"], ["a==1"]),
+        ({PoetryPackage("a", "1"), PoetryPackage("b")}, ["a<=2"], ["a==1"]),
+        ({PoetryPackage("a", "1"), PoetryPackage("b")}, ["a>=1"], ["a==1"]),
     ],
 )
-def test__sync_hooks_additional_deps__no_new_deps(
+def test__sync_hooks_additional_dependencies__no_new_deps(
     poetry_deps, additional_deps, expected_additional_deps
 ) -> None:
     """Check that `_sync_hooks_additional_dependencies` handles the different ways to write a package entry."""