From 0d93ab11610a5c36310b708d5a69d098b2606f7d Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Tue, 28 Nov 2023 09:52:09 +0100
Subject: [PATCH 01/10] ENH: type Python version with `Literal`

---
 src/repoma/utilities/project_info.py | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/src/repoma/utilities/project_info.py b/src/repoma/utilities/project_info.py
index 66ecf712..4729a3ab 100644
--- a/src/repoma/utilities/project_info.py
+++ b/src/repoma/utilities/project_info.py
@@ -1,6 +1,7 @@
 """Helper functions for reading from and writing to :file:`setup.cfg`."""
 
 import os
+import sys
 from configparser import ConfigParser
 from textwrap import dedent
 from typing import Dict, List, Optional
@@ -14,11 +15,19 @@
 from . import CONFIG_PATH
 from .cfg import open_config
 
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
+
+PythonVersion = Literal["3.6", "3.7", "3.8", "3.9", "3.10", "3.11", "3.12"]
+
 
 @frozen
 class ProjectInfo:
     name: Optional[str] = None
-    supported_python_versions: Optional[List[str]] = None
+    supported_python_versions: Optional[List[PythonVersion]] = None
     urls: Dict[str, str] = field(factory=dict)
 
     def is_empty(self) -> bool:
@@ -80,13 +89,13 @@ def get_project_info(pyproject: Optional[TOMLDocument] = None) -> ProjectInfo:
     raise PrecommitError(msg)
 
 
-def _extract_python_versions(classifiers: List[str]) -> Optional[List[str]]:
+def _extract_python_versions(classifiers: List[str]) -> Optional[List[PythonVersion]]:
     identifier = "Programming Language :: Python :: 3."
     version_classifiers = [s for s in classifiers if s.startswith(identifier)]
     if not version_classifiers:
         return None
     prefix = identifier[:-2]
-    return [s.replace(prefix, "") for s in version_classifiers]
+    return [s.replace(prefix, "") for s in version_classifiers]  # type: ignore[misc]
 
 
 def get_pypi_name(pyproject: Optional[TOMLDocument] = None) -> str:
@@ -107,7 +116,7 @@ def get_pypi_name(pyproject: Optional[TOMLDocument] = None) -> str:
 
 def get_supported_python_versions(
     pyproject: Optional[TOMLDocument] = None,
-) -> List[str]:
+) -> List[PythonVersion]:
     """Extract supported Python versions from package classifiers.
 
     >>> get_supported_python_versions()

From 25ba52dc21d4c93f85f6bdda64c5415dead6b869 Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Tue, 28 Nov 2023 10:44:46 +0100
Subject: [PATCH 02/10] MAINT: extract `add_dependency()` function for
 `pyproject.toml`

---
 src/repoma/check_dev_files/ruff.py | 35 +++++-------------------------
 src/repoma/utilities/pyproject.py  | 32 ++++++++++++++++++++++++++-
 2 files changed, 37 insertions(+), 30 deletions(-)

diff --git a/src/repoma/check_dev_files/ruff.py b/src/repoma/check_dev_files/ruff.py
index c966ceaa..2e096e32 100644
--- a/src/repoma/check_dev_files/ruff.py
+++ b/src/repoma/check_dev_files/ruff.py
@@ -1,7 +1,6 @@
 """Check `Ruff <https://ruff.rs>`_ configuration."""
 
 import os
-from copy import deepcopy
 from textwrap import dedent
 from typing import List, Set
 
@@ -27,6 +26,7 @@
     open_setup_cfg,
 )
 from repoma.utilities.pyproject import (
+    add_dependency,
     complies_with_subset,
     get_sub_table,
     load_pyproject,
@@ -230,39 +230,16 @@ def _update_pyproject() -> None:
             f" [{CONFIG_PATH.pyproject}]"
         )
         raise PrecommitError(msg)
-    project = get_sub_table(pyproject, "project", create=True)
-    old_dependencies = project.get("optional-dependencies")
-    new_dependencies = deepcopy(old_dependencies)
     python_versions = project_info.supported_python_versions
     if python_versions is not None and "3.6" in python_versions:
         ruff = 'ruff; python_version >="3.7.0"'
     else:
         ruff = "ruff"
-    if new_dependencies is None:
-        new_dependencies = dict(
-            dev=[f"{package}[sty]"],
-            lint=[ruff],
-            sty=[f"{package}[lint]"],
-        )
-    else:
-        __add_package(new_dependencies, "dev", f"{package}[sty]")
-        __add_package(new_dependencies, "lint", ruff)
-        __add_package(new_dependencies, "sty", f"{package}[lint]")
-    if old_dependencies != new_dependencies:
-        project["optional-dependencies"] = new_dependencies
-        write_pyproject(pyproject)
-        msg = f"Updated [project.optional-dependencies] in {CONFIG_PATH.pyproject}"
-        raise PrecommitError(msg)
-
-
-def __add_package(optional_dependencies: Table, key: str, package: str) -> None:
-    section = optional_dependencies.get(key)
-    if section is None:
-        optional_dependencies[key] = [package]
-    elif package not in section:
-        optional_dependencies[key] = to_toml_array(
-            sorted({package, *section}, key=lambda s: ('"' in s, s))  # Taplo sorting
-        )
+    executor = Executor()
+    executor(add_dependency, ruff, key="lint")
+    executor(add_dependency, f"{package}[lint]", key="sty")
+    executor(add_dependency, f"{package}[sty]", key="dev")
+    executor.finalize()
 
 
 def _remove_nbqa() -> None:
diff --git a/src/repoma/utilities/pyproject.py b/src/repoma/utilities/pyproject.py
index 488272b1..f1e09b75 100644
--- a/src/repoma/utilities/pyproject.py
+++ b/src/repoma/utilities/pyproject.py
@@ -1,7 +1,7 @@
 """Tools for loading, inspecting, and updating :code:`pyproject.toml`."""
 
 import os
-from typing import Any, Iterable, Optional
+from typing import Any, Iterable, List, Optional, Set
 
 import tomlkit
 from tomlkit.container import Container
@@ -13,6 +13,36 @@
 from repoma.utilities.precommit import find_repo, load_round_trip_precommit_config
 
 
+def add_dependency(package: str, optional_key: Optional[str] = None) -> None:
+    pyproject = load_pyproject()
+    if optional_key is None:
+        project = get_sub_table(pyproject, "project", create=True)
+        existing_dependencies: Set[str] = set(project.get("dependencies", []))
+        if package in existing_dependencies:
+            return
+        existing_dependencies.add(package)
+        project["dependencies"] = to_toml_array(_sort_taplo(existing_dependencies))
+    else:
+        optional_dependencies = get_sub_table(
+            pyproject, "project.optional-dependencies", create=True
+        )
+        existing_dependencies = set(optional_dependencies.get(optional_key, []))
+        if package in existing_dependencies:
+            return
+        existing_dependencies.add(package)
+        existing_dependencies = set(existing_dependencies)
+        optional_dependencies[optional_key] = to_toml_array(
+            _sort_taplo(existing_dependencies)
+        )
+    write_pyproject(pyproject)
+    msg = f"Listed {package} as a dependency under {CONFIG_PATH.pyproject}"
+    raise PrecommitError(msg)
+
+
+def _sort_taplo(items: Iterable[str]) -> List[str]:
+    return sorted(items, key=lambda s: ('"' in s, s))
+
+
 def complies_with_subset(settings: dict, minimal_settings: dict) -> bool:
     return all(settings.get(key) == value for key, value in minimal_settings.items())
 

From 4b680781470c5a52cb758348fb198c9646740aa9 Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Tue, 28 Nov 2023 11:23:17 +0100
Subject: [PATCH 03/10] MAINT: allow nested keys in `add_dependency()`

---
 src/repoma/check_dev_files/ruff.py |  6 +----
 src/repoma/utilities/pyproject.py  | 43 +++++++++++++++++++++++++++---
 2 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/src/repoma/check_dev_files/ruff.py b/src/repoma/check_dev_files/ruff.py
index 2e096e32..65cee492 100644
--- a/src/repoma/check_dev_files/ruff.py
+++ b/src/repoma/check_dev_files/ruff.py
@@ -235,11 +235,7 @@ def _update_pyproject() -> None:
         ruff = 'ruff; python_version >="3.7.0"'
     else:
         ruff = "ruff"
-    executor = Executor()
-    executor(add_dependency, ruff, key="lint")
-    executor(add_dependency, f"{package}[lint]", key="sty")
-    executor(add_dependency, f"{package}[sty]", key="dev")
-    executor.finalize()
+    add_dependency(ruff, optional_key=["lint", "sty", "dev"])
 
 
 def _remove_nbqa() -> None:
diff --git a/src/repoma/utilities/pyproject.py b/src/repoma/utilities/pyproject.py
index f1e09b75..37c2b968 100644
--- a/src/repoma/utilities/pyproject.py
+++ b/src/repoma/utilities/pyproject.py
@@ -1,7 +1,9 @@
 """Tools for loading, inspecting, and updating :code:`pyproject.toml`."""
 
 import os
-from typing import Any, Iterable, List, Optional, Set
+from collections import abc
+from itertools import zip_longest
+from typing import Any, Iterable, List, Optional, Sequence, Set, Union
 
 import tomlkit
 from tomlkit.container import Container
@@ -10,10 +12,13 @@
 
 from repoma.errors import PrecommitError
 from repoma.utilities import CONFIG_PATH
+from repoma.utilities.executor import Executor
 from repoma.utilities.precommit import find_repo, load_round_trip_precommit_config
 
 
-def add_dependency(package: str, optional_key: Optional[str] = None) -> None:
+def add_dependency(
+    package: str, optional_key: Optional[Union[str, Sequence[str]]] = None
+) -> None:
     pyproject = load_pyproject()
     if optional_key is None:
         project = get_sub_table(pyproject, "project", create=True)
@@ -22,7 +27,7 @@ def add_dependency(package: str, optional_key: Optional[str] = None) -> None:
             return
         existing_dependencies.add(package)
         project["dependencies"] = to_toml_array(_sort_taplo(existing_dependencies))
-    else:
+    elif isinstance(optional_key, str):
         optional_dependencies = get_sub_table(
             pyproject, "project.optional-dependencies", create=True
         )
@@ -34,6 +39,18 @@ def add_dependency(package: str, optional_key: Optional[str] = None) -> None:
         optional_dependencies[optional_key] = to_toml_array(
             _sort_taplo(existing_dependencies)
         )
+    elif isinstance(optional_key, abc.Iterable):
+        this_package = get_package_name_safe(pyproject)
+        executor = Executor()
+        for key, next_key in zip_longest(optional_key, optional_key[1:]):
+            if next_key is None:
+                executor(add_dependency, package, key)
+            else:
+                executor(add_dependency, f"{this_package}[{key}]", next_key)
+        executor.finalize()
+    else:
+        msg = f"Unsupported type for optional_key: {type(optional_key)}"
+        raise NotImplementedError(msg)
     write_pyproject(pyproject)
     msg = f"Listed {package} as a dependency under {CONFIG_PATH.pyproject}"
     raise PrecommitError(msg)
@@ -56,6 +73,26 @@ def load_pyproject(content: Optional[str] = None) -> TOMLDocument:
     return tomlkit.loads(content)
 
 
+def get_package_name(pyproject: Optional[TOMLDocument]) -> Optional[str]:
+    pyproject = load_pyproject()
+    project = get_sub_table(pyproject, "project")
+    package_name = project.get("name")
+    if package_name is None:
+        return None
+    return package_name
+
+
+def get_package_name_safe(pyproject: Optional[TOMLDocument]) -> str:
+    package_name = get_package_name(pyproject)
+    if package_name is None:
+        msg = (
+            "Please specify a [project.name] for the package in"
+            f" [{CONFIG_PATH.pyproject}]"
+        )
+        raise PrecommitError(msg)
+    return package_name
+
+
 def get_sub_table(config: Container, dotted_header: str, create: bool = False) -> Table:
     """Get a TOML sub-table through a dotted header key."""
     current_table: Any = config

From 261dc0e0fc581ae2421113423d55597e23b416ef Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Tue, 28 Nov 2023 12:52:05 +0100
Subject: [PATCH 04/10] MAINT: load pyproject from stream

---
 src/repoma/utilities/pyproject.py | 24 ++++++++++------
 tests/utilities/test_pyproject.py | 47 +++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 9 deletions(-)
 create mode 100644 tests/utilities/test_pyproject.py

diff --git a/src/repoma/utilities/pyproject.py b/src/repoma/utilities/pyproject.py
index 37c2b968..a6a54733 100644
--- a/src/repoma/utilities/pyproject.py
+++ b/src/repoma/utilities/pyproject.py
@@ -1,9 +1,10 @@
 """Tools for loading, inspecting, and updating :code:`pyproject.toml`."""
 
-import os
+import io
 from collections import abc
 from itertools import zip_longest
-from typing import Any, Iterable, List, Optional, Sequence, Set, Union
+from pathlib import Path
+from typing import IO, Any, Iterable, List, Optional, Sequence, Set, Union
 
 import tomlkit
 from tomlkit.container import Container
@@ -64,13 +65,18 @@ def complies_with_subset(settings: dict, minimal_settings: dict) -> bool:
     return all(settings.get(key) == value for key, value in minimal_settings.items())
 
 
-def load_pyproject(content: Optional[str] = None) -> TOMLDocument:
-    if not os.path.exists(CONFIG_PATH.pyproject):
-        return TOMLDocument()
-    if content is None:
-        with open(CONFIG_PATH.pyproject) as stream:
-            return tomlkit.loads(stream.read())
-    return tomlkit.loads(content)
+def load_pyproject(
+    source: Union[IO, Path, str] = CONFIG_PATH.pyproject
+) -> TOMLDocument:
+    if isinstance(source, io.IOBase):
+        return tomlkit.load(source)
+    if isinstance(source, Path):
+        with open(source) as stream:
+            return load_pyproject(stream)
+    if isinstance(source, str):
+        return tomlkit.loads(source)
+    msg = f"Source of type {type(source).__name__} is not supported"
+    raise TypeError(msg)
 
 
 def get_package_name(pyproject: Optional[TOMLDocument]) -> Optional[str]:
diff --git a/tests/utilities/test_pyproject.py b/tests/utilities/test_pyproject.py
new file mode 100644
index 00000000..68c3e26d
--- /dev/null
+++ b/tests/utilities/test_pyproject.py
@@ -0,0 +1,47 @@
+from pathlib import Path
+from textwrap import dedent
+from typing import Optional
+
+import pytest
+from tomlkit.items import Table
+
+from repoma.utilities.pyproject import load_pyproject
+
+REPOMA_DIR = Path(__file__).absolute().parent.parent.parent
+
+
+@pytest.mark.parametrize("path", [None, REPOMA_DIR / "pyproject.toml"])
+def test_load_pyproject(path: Optional[Path]):
+    if path is None:
+        pyproject = load_pyproject()
+    else:
+        pyproject = load_pyproject(path)
+    assert "build-system" in pyproject
+    assert "tool" in pyproject
+
+
+def test_load_pyproject_str():
+    src = dedent("""
+    [build-system]
+    build-backend = "setuptools.build_meta"
+    requires = [
+        "setuptools>=61.2",
+        "setuptools_scm",
+    ]
+
+    [project]
+    dependencies = [
+        "attrs",
+        "sympy >=1.10",
+    ]
+    name = "my-package"
+    requires-python = ">=3.7"
+    """)
+    pyproject = load_pyproject(src)
+    assert isinstance(pyproject["build-system"], Table)
+    assert pyproject["project"]["dependencies"] == ["attrs", "sympy >=1.10"]  # type: ignore[index]
+
+
+def test_load_pyproject_type_error():
+    with pytest.raises(TypeError, match="Source of type int is not supported"):
+        _ = load_pyproject(1)  # type: ignore[arg-type]

From 8c7c12c467b0809bf3fa8034ad59dfdac2e55ec3 Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Wed, 29 Nov 2023 11:42:53 +0100
Subject: [PATCH 05/10] ENH: add target argument for `write_pyproject()`

---
 src/repoma/utilities/pyproject.py | 16 +++++++---
 tests/utilities/test_pyproject.py | 52 +++++++++++++++++++++++++++++--
 2 files changed, 62 insertions(+), 6 deletions(-)

diff --git a/src/repoma/utilities/pyproject.py b/src/repoma/utilities/pyproject.py
index a6a54733..17dc9355 100644
--- a/src/repoma/utilities/pyproject.py
+++ b/src/repoma/utilities/pyproject.py
@@ -113,10 +113,18 @@ def get_sub_table(config: Container, dotted_header: str, create: bool = False) -
     return current_table
 
 
-def write_pyproject(config: TOMLDocument) -> None:
-    src = tomlkit.dumps(config, sort_keys=True)
-    with open(CONFIG_PATH.pyproject, "w") as stream:
-        stream.write(src)
+def write_pyproject(
+    config: TOMLDocument, target: Union[IO, Path, str] = CONFIG_PATH.pyproject
+) -> None:
+    if isinstance(target, io.IOBase):
+        tomlkit.dump(config, target, sort_keys=True)
+    elif isinstance(target, (Path, str)):
+        src = tomlkit.dumps(config, sort_keys=True)
+        with open(target, "w") as stream:
+            stream.write(src)
+    else:
+        msg = f"Target of type {type(target).__name__} is not supported"
+        raise TypeError(msg)
 
 
 def to_toml_array(items: Iterable[Any], enforce_multiline: bool = False) -> Array:
diff --git a/tests/utilities/test_pyproject.py b/tests/utilities/test_pyproject.py
index 68c3e26d..fef336ed 100644
--- a/tests/utilities/test_pyproject.py
+++ b/tests/utilities/test_pyproject.py
@@ -1,15 +1,63 @@
+import io
 from pathlib import Path
-from textwrap import dedent
+from textwrap import dedent, indent
 from typing import Optional
 
 import pytest
 from tomlkit.items import Table
 
-from repoma.utilities.pyproject import load_pyproject
+from repoma.utilities.pyproject import (
+    get_sub_table,
+    load_pyproject,
+    to_toml_array,
+    write_pyproject,
+)
 
 REPOMA_DIR = Path(__file__).absolute().parent.parent.parent
 
 
+def test_edit_toml():
+    src = dedent("""
+    [owner]
+    name = "John Smith"
+    age = 30
+
+    [owner.address]
+    city = "Wonderland"
+    street = "123 Main St"
+    """)
+    config = load_pyproject(src)
+
+    address = get_sub_table(config, "owner.address")
+    address["city"] = "New York"
+    work = get_sub_table(config, "owner.work", create=True)
+    work["type"] = "scientist"
+    tools = get_sub_table(config, "tool", create=True)
+    tools["black"] = to_toml_array(["--line-length=79"], enforce_multiline=True)
+
+    stream = io.StringIO()
+    write_pyproject(config, target=stream)
+    result = stream.getvalue()
+    print(indent(result, prefix=4 * " "))  # noqa: T201  # run with pytest -s
+    assert result == dedent("""
+    [owner]
+    name = "John Smith"
+    age = 30
+
+    [owner.address]
+    city = "New York"
+    street = "123 Main St"
+
+    [owner.work]
+    type = "scientist"
+
+    [tool]
+    black = [
+        "--line-length=79",
+    ]
+    """)
+
+
 @pytest.mark.parametrize("path", [None, REPOMA_DIR / "pyproject.toml"])
 def test_load_pyproject(path: Optional[Path]):
     if path is None:

From 30ebfa67a7a2cb707f8ddd97e876747166fbcd68 Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Wed, 29 Nov 2023 17:03:06 +0100
Subject: [PATCH 06/10] ENH!: reset stream after before loading/writing

---
 src/repoma/utilities/pyproject.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/repoma/utilities/pyproject.py b/src/repoma/utilities/pyproject.py
index 17dc9355..b3fbcf64 100644
--- a/src/repoma/utilities/pyproject.py
+++ b/src/repoma/utilities/pyproject.py
@@ -69,6 +69,7 @@ def load_pyproject(
     source: Union[IO, Path, str] = CONFIG_PATH.pyproject
 ) -> TOMLDocument:
     if isinstance(source, io.IOBase):
+        source.seek(0)
         return tomlkit.load(source)
     if isinstance(source, Path):
         with open(source) as stream:
@@ -117,6 +118,7 @@ def write_pyproject(
     config: TOMLDocument, target: Union[IO, Path, str] = CONFIG_PATH.pyproject
 ) -> None:
     if isinstance(target, io.IOBase):
+        target.seek(0)
         tomlkit.dump(config, target, sort_keys=True)
     elif isinstance(target, (Path, str)):
         src = tomlkit.dumps(config, sort_keys=True)

From a77fca7bf293773a3ffbaf7d7e78335625844cd7 Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Wed, 29 Nov 2023 17:09:44 +0100
Subject: [PATCH 07/10] MAINT: test `get_package_name()` function

---
 src/repoma/utilities/pyproject.py | 21 ++++++++++++++-------
 tests/utilities/test_pyproject.py | 15 +++++++++++++++
 2 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/src/repoma/utilities/pyproject.py b/src/repoma/utilities/pyproject.py
index b3fbcf64..467ef755 100644
--- a/src/repoma/utilities/pyproject.py
+++ b/src/repoma/utilities/pyproject.py
@@ -80,21 +80,28 @@ def load_pyproject(
     raise TypeError(msg)
 
 
-def get_package_name(pyproject: Optional[TOMLDocument]) -> Optional[str]:
-    pyproject = load_pyproject()
-    project = get_sub_table(pyproject, "project")
+def get_package_name(
+    source: Union[IO, Path, TOMLDocument, str] = CONFIG_PATH.pyproject
+) -> Optional[str]:
+    if isinstance(source, TOMLDocument):
+        pyproject = source
+    else:
+        pyproject = load_pyproject(source)
+    project = get_sub_table(pyproject, "project", create=True)
     package_name = project.get("name")
     if package_name is None:
         return None
     return package_name
 
 
-def get_package_name_safe(pyproject: Optional[TOMLDocument]) -> str:
-    package_name = get_package_name(pyproject)
+def get_package_name_safe(
+    source: Union[IO, Path, TOMLDocument, str] = CONFIG_PATH.pyproject
+) -> str:
+    package_name = get_package_name(source)
     if package_name is None:
         msg = (
-            "Please specify a [project.name] for the package in"
-            f" [{CONFIG_PATH.pyproject}]"
+            "Please provide a name for the package under the [project] table in"
+            f" {CONFIG_PATH.pyproject}"
         )
         raise PrecommitError(msg)
     return package_name
diff --git a/tests/utilities/test_pyproject.py b/tests/utilities/test_pyproject.py
index fef336ed..faead1fe 100644
--- a/tests/utilities/test_pyproject.py
+++ b/tests/utilities/test_pyproject.py
@@ -6,7 +6,9 @@
 import pytest
 from tomlkit.items import Table
 
+from repoma.errors import PrecommitError
 from repoma.utilities.pyproject import (
+    get_package_name_safe,
     get_sub_table,
     load_pyproject,
     to_toml_array,
@@ -58,6 +60,19 @@ def test_edit_toml():
     """)
 
 
+def test_get_package_name_safe():
+    correct_input = io.StringIO(dedent("""
+    [project]
+    name = "my-package"
+    """))
+    assert get_package_name_safe(correct_input) == "my-package"
+
+    with pytest.raises(PrecommitError, match=r"^Please provide a name for the package"):
+        _ = get_package_name_safe(io.StringIO("[project]"))
+    with pytest.raises(PrecommitError, match=r"^Please provide a name for the package"):
+        _ = get_package_name_safe(io.StringIO())
+
+
 @pytest.mark.parametrize("path", [None, REPOMA_DIR / "pyproject.toml"])
 def test_load_pyproject(path: Optional[Path]):
     if path is None:

From 2be67a6dba915d2ac73fd9c03fa7abc990eb3d10 Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Wed, 29 Nov 2023 17:10:46 +0100
Subject: [PATCH 08/10] MAINT: test `add_dependency()` function

---
 src/repoma/utilities/pyproject.py | 35 ++++++++++++-----
 tests/utilities/test_pyproject.py | 62 +++++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+), 10 deletions(-)

diff --git a/src/repoma/utilities/pyproject.py b/src/repoma/utilities/pyproject.py
index 467ef755..97e61247 100644
--- a/src/repoma/utilities/pyproject.py
+++ b/src/repoma/utilities/pyproject.py
@@ -2,7 +2,6 @@
 
 import io
 from collections import abc
-from itertools import zip_longest
 from pathlib import Path
 from typing import IO, Any, Iterable, List, Optional, Sequence, Set, Union
 
@@ -17,10 +16,21 @@
 from repoma.utilities.precommit import find_repo, load_round_trip_precommit_config
 
 
-def add_dependency(
-    package: str, optional_key: Optional[Union[str, Sequence[str]]] = None
+def add_dependency(  # noqa: C901, PLR0912
+    package: str,
+    optional_key: Optional[Union[str, Sequence[str]]] = None,
+    source: Union[IO, Path, TOMLDocument, str] = CONFIG_PATH.pyproject,
+    target: Optional[Union[IO, Path, str]] = None,
 ) -> None:
-    pyproject = load_pyproject()
+    if isinstance(source, TOMLDocument):
+        pyproject = source
+    else:
+        pyproject = load_pyproject(source)
+    if target is None:
+        if isinstance(source, TOMLDocument):
+            msg = "If the source is a TOML document, you have to specify a target"
+            raise TypeError(msg)
+        target = source
     if optional_key is None:
         project = get_sub_table(pyproject, "project", create=True)
         existing_dependencies: Set[str] = set(project.get("dependencies", []))
@@ -40,19 +50,24 @@ def add_dependency(
         optional_dependencies[optional_key] = to_toml_array(
             _sort_taplo(existing_dependencies)
         )
-    elif isinstance(optional_key, abc.Iterable):
+    elif isinstance(optional_key, abc.Sequence):
+        if len(optional_key) < 2:  # noqa: PLR2004
+            msg = "Need at least two keys to define nested optional dependencies"
+            raise ValueError(msg)
         this_package = get_package_name_safe(pyproject)
         executor = Executor()
-        for key, next_key in zip_longest(optional_key, optional_key[1:]):
-            if next_key is None:
-                executor(add_dependency, package, key)
+        for key, previous in zip(optional_key, [None, *optional_key]):
+            if previous is None:
+                executor(add_dependency, package, key, source, target)
             else:
-                executor(add_dependency, f"{this_package}[{key}]", next_key)
+                executor(
+                    add_dependency, f"{this_package}[{previous}]", key, source, target
+                )
         executor.finalize()
     else:
         msg = f"Unsupported type for optional_key: {type(optional_key)}"
         raise NotImplementedError(msg)
-    write_pyproject(pyproject)
+    write_pyproject(pyproject, target)
     msg = f"Listed {package} as a dependency under {CONFIG_PATH.pyproject}"
     raise PrecommitError(msg)
 
diff --git a/tests/utilities/test_pyproject.py b/tests/utilities/test_pyproject.py
index faead1fe..48656af0 100644
--- a/tests/utilities/test_pyproject.py
+++ b/tests/utilities/test_pyproject.py
@@ -8,6 +8,7 @@
 
 from repoma.errors import PrecommitError
 from repoma.utilities.pyproject import (
+    add_dependency,
     get_package_name_safe,
     get_sub_table,
     load_pyproject,
@@ -18,6 +19,67 @@
 REPOMA_DIR = Path(__file__).absolute().parent.parent.parent
 
 
+def test_add_dependency():
+    stream = io.StringIO(dedent("""
+    [project]
+    name = "my-package"
+    """))
+    stream.seek(0)
+    dependency = "attrs"
+    with pytest.raises(
+        PrecommitError,
+        match=f"Listed {dependency} as a dependency under pyproject.toml",
+    ):
+        add_dependency(dependency, source=stream)
+    result = stream.getvalue()
+    print(result)  # noqa: T201  # run with pytest -s
+    assert result == dedent("""
+    [project]
+    name = "my-package"
+    dependencies = ["attrs"]
+    """)
+
+
+def test_add_dependency_nested():
+    stream = io.StringIO(dedent("""
+    [project]
+    name = "my-package"
+    """))
+    stream.seek(0)
+    with pytest.raises(PrecommitError):
+        add_dependency("ruff", optional_key=["lint", "sty", "dev"], source=stream)
+    result = stream.getvalue()
+    print(result)  # noqa: T201  # run with pytest -s
+    assert result == dedent("""
+    [project]
+    name = "my-package"
+
+    [project.optional-dependencies]
+    lint = ["ruff"]
+    sty = ["my-package[lint]"]
+    dev = ["my-package[sty]"]
+    """)
+
+
+def test_add_dependency_optional():
+    stream = io.StringIO(dedent("""
+    [project]
+    name = "my-package"
+    """))
+    stream.seek(0)
+    with pytest.raises(PrecommitError):
+        add_dependency("ruff", optional_key="lint", source=stream)
+    result = stream.getvalue()
+    print(result)  # noqa: T201  # run with pytest -s
+    assert result == dedent("""
+    [project]
+    name = "my-package"
+
+    [project.optional-dependencies]
+    lint = ["ruff"]
+    """)
+
+
 def test_edit_toml():
     src = dedent("""
     [owner]

From e2c3ddfe2b4289e68b275c13b60700bf8f6e5225 Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Wed, 29 Nov 2023 17:18:42 +0100
Subject: [PATCH 09/10] DX: automatically list `jupyterlab` requirements

---
 .cspell.json                           |  1 +
 src/repoma/check_dev_files/__init__.py |  3 +++
 src/repoma/check_dev_files/jupyter.py  | 26 ++++++++++++++++++++++++++
 src/repoma/utilities/pyproject.py      |  3 ++-
 4 files changed, 32 insertions(+), 1 deletion(-)
 create mode 100644 src/repoma/check_dev_files/jupyter.py

diff --git a/.cspell.json b/.cspell.json
index 72da7eb2..dfb23319 100644
--- a/.cspell.json
+++ b/.cspell.json
@@ -71,6 +71,7 @@
     "indentless",
     "ipynb",
     "jsonschema",
+    "jupyterlab",
     "linkcheck",
     "maxdepth",
     "maxsplit",
diff --git a/src/repoma/check_dev_files/__init__.py b/src/repoma/check_dev_files/__init__.py
index 1cfdb2aa..25f3748b 100644
--- a/src/repoma/check_dev_files/__init__.py
+++ b/src/repoma/check_dev_files/__init__.py
@@ -16,6 +16,7 @@
     github_labels,
     github_workflows,
     gitpod,
+    jupyter,
     mypy,
     nbstripout,
     precommit,
@@ -58,6 +59,8 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
             skip_tests=_to_list(args.ci_skipped_tests),
             test_extras=_to_list(args.ci_test_extras),
         )
+    if has_notebooks:
+        executor(jupyter.main)
     executor(nbstripout.main)
     executor(toml.main)  # has to run before pre-commit
     executor(prettier.main, args.no_prettierrc)
diff --git a/src/repoma/check_dev_files/jupyter.py b/src/repoma/check_dev_files/jupyter.py
new file mode 100644
index 00000000..9fa4dd93
--- /dev/null
+++ b/src/repoma/check_dev_files/jupyter.py
@@ -0,0 +1,26 @@
+"""Update the developer setup when using Jupyter notebooks."""
+
+from repoma.utilities.executor import Executor
+from repoma.utilities.project_info import get_supported_python_versions
+from repoma.utilities.pyproject import add_dependency
+
+
+def main() -> None:
+    _update_dev_requirements()
+
+
+def _update_dev_requirements() -> None:
+    if "3.6" in get_supported_python_versions():
+        return
+    hierarchy = ["jupyter", "dev"]
+    dependencies = [
+        "jupyterlab",
+        "jupyterlab-code-formatter",
+        "jupyterlab-lsp",
+        "jupyterlab-myst",
+        "python-lsp-server[rope]",
+    ]
+    executor = Executor()
+    for dependency in dependencies:
+        executor(add_dependency, dependency, optional_key=hierarchy)
+    executor.finalize()
diff --git a/src/repoma/utilities/pyproject.py b/src/repoma/utilities/pyproject.py
index 97e61247..77ef223e 100644
--- a/src/repoma/utilities/pyproject.py
+++ b/src/repoma/utilities/pyproject.py
@@ -63,7 +63,8 @@ def add_dependency(  # noqa: C901, PLR0912
                 executor(
                     add_dependency, f"{this_package}[{previous}]", key, source, target
                 )
-        executor.finalize()
+        if executor.finalize() == 0:
+            return
     else:
         msg = f"Unsupported type for optional_key: {type(optional_key)}"
         raise NotImplementedError(msg)

From 18e3bee9c1f9db38208595712eef84758160e998 Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Wed, 29 Nov 2023 18:28:45 +0100
Subject: [PATCH 10/10] ENH: remove final newline

---
 src/repoma/utilities/pyproject.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/repoma/utilities/pyproject.py b/src/repoma/utilities/pyproject.py
index 77ef223e..1a3431ea 100644
--- a/src/repoma/utilities/pyproject.py
+++ b/src/repoma/utilities/pyproject.py
@@ -145,6 +145,7 @@ def write_pyproject(
         tomlkit.dump(config, target, sort_keys=True)
     elif isinstance(target, (Path, str)):
         src = tomlkit.dumps(config, sort_keys=True)
+        src = f"{src.strip()}\n"
         with open(target, "w") as stream:
             stream.write(src)
     else: