feat: csv shifting

nuh_helper/book_page/__init__.py

@@ -0,0 +1,11 @@
+from nuh_helper.book_page.csv_book import CSVBook
+from nuh_helper.book_page.generic import Book, Page, book_open
+from nuh_helper.book_page.xlsx_book import XLSXBook
+
+__all__ = [  # noqa: PLE0604
+    CSVBook,
+    XLSXBook,
+    Book,
+    Page,
+    book_open,
+]

nuh_helper/book_page/book.py

@@ -0,0 +1,37 @@
+from abc import ABC, abstractmethod
+from collections.abc import Generator as gen
+from collections.abc import Iterable as iterator
+from functools import cached_property
+from pathlib import Path
+
+
+class Page(ABC):
+    def __init__(self, name: str, page: object) -> None:
+        self._name = name
+        self._page = page
+
+    @abstractmethod
+    def stream_rows(self) -> gen[tuple[int, list[str]]]: ...
+
+    @cached_property
+    def name(self) -> str:
+        return self._name
+
+
+class Book(ABC):
+    def __contains__(self, name: str) -> bool:
+        return name in self.keys()
+
+    @abstractmethod
+    def __getitem__(self, name: str) -> gen[Page]: ...
+
+    @abstractmethod
+    def keys(self) -> iterator[str]: ...
+
+    @abstractmethod
+    def stream_copy(self, into: Path) -> gen[tuple[str, int, list[str]]]: ...
+
+    def __iter__(self) -> gen[Page]:
+        for name in self.keys():
+            with self[name] as page:
+                yield page

nuh_helper/book_page/csv_book.py

@@ -0,0 +1,114 @@
+import contextlib
+import csv
+from _collections_abc import dict_keys
+from collections.abc import Generator as gen
+from pathlib import Path
+
+from nuh_helper.book_page.book import Book, Page
+
+
+class CSVBook(Book):
+    def __init__(self, path: str | Path | list[Path]) -> None:
+        self._page_files = {page.stem.lower(): page for page in csv_files(path)}
+
+    def keys(self) -> dict_keys:
+        return self._page_files.keys()
+
+    @contextlib.contextmanager
+    def __getitem__(self, name: str) -> gen[Page]:
+
+        class CSVPage(Page):
+            def stream_rows(self) -> gen[tuple[int, list[str]]]:
+                assert isinstance(self._page, Path)
+                with self._page.open() as file:
+                    for row, line in enumerate(
+                        csv.reader(
+                            file,
+                            delimiter="\t" if file.name.endswith(".tsv") else ",",
+                        )
+                    ):
+                        # strip cell strings
+                        line = [
+                            cell.strip() if isinstance(cell, str) else cell
+                            for cell in line
+                        ]
+
+                        # convert empty cells to None
+                        line = [cell if cell else None for cell in line]
+
+                        # strip bom
+                        if (
+                            row == 0
+                            and isinstance(line[0], str)
+                            and line[0][:1] == "\ufeff"
+                        ):
+                            line[0] = line[0][1:]
+
+                        yield (
+                            row,
+                            line,
+                        )
+
+        assert isinstance(self._page_files[name], Path)
+        yield CSVPage(name, self._page_files[name])
+
+    def stream_copy(self, into: str | Path) -> gen[tuple[str, int, list[str]]]:
+        if not isinstance(into, Path):
+            into = Path(into)
+
+        if into.is_file():
+            raise RuntimeError(f"can't write csvs to {into} because it's a file")
+
+        into.mkdir(parents=True, exist_ok=True)
+
+        for page in self:
+            name = page.name
+            file = self._page_files[name].name
+            with (into / file).open("w") as data:
+                writer = csv.writer(
+                    data, delimiter="\t" if file.endswith(".tsv") else ","
+                )
+
+                for row, cell in page.stream_rows():
+                    yield name, row, cell
+                    writer.writerow(cell)
+
+
+def csv_files(path: str | Path | list[Path]) -> list[Path]:
+    """process whatever form the path takes into a list of Path objects"""
+    if isinstance(path, str):
+        path = Path(path)
+
+    # get it to a list
+    if isinstance(path, Path):
+        if path.is_file():
+            raise RuntimeError(f"need a dir or list of files, but, got a file {path=}")
+        path = sorted(
+            list(path.glob("*.csv")) + list(path.glob("*.tsv")),
+            key=lambda file: file.stem.lower(),
+        )
+
+    # compute dupes
+    collisions = sorted(
+        [
+            page
+            for page in path
+            if len([them for them in path if them.stem.lower() == page.stem.lower()])
+            != 1
+        ],
+        key=lambda file: file.name.lower(),
+    )
+    # raise error if there are dupes
+    if collisions:
+        raise ValueError(
+            f"the following csv files have name collisions\n    {
+                '\n    '.join(
+                    [
+                        (file.stem.lower() + '\n        ' + str(file))
+                        for file in collisions
+                    ]
+                )
+            }"
+        )
+
+    return path

nuh_helper/book_page/generic.py

@@ -0,0 +1,25 @@
+from pathlib import Path
+
+from .book import Page
+from .csv_book import CSVBook
+from .xlsx_book import XLSXBook
+
+# use this or-type instead of the "real" base class since we only have/support these two
+Book = CSVBook | XLSXBook
+
+
+def book_open(src: str | Path | list[Path]) -> Book:
+    if isinstance(src, str):
+        src = Path(src)
+
+    if isinstance(src, Path) and src.name.endswith(".xlsx"):
+        return XLSXBook(src)
+    else:
+        return CSVBook(src)
+
+
+__all__ = [  # noqa: PLE0604
+    Book,
+    Page,
+    book_open,
+]

nuh_helper/book_page/xlsx_book.py

@@ -0,0 +1,61 @@
+import contextlib
+import shutil
+from collections.abc import Generator as gen
+from pathlib import Path
+
+import openpyxl
+
+from .book import Book, Page
+
+
+class XLSXBook(Book):
+    def __init__(self, file: str | Path, read_only: bool = True) -> None:
+        if isinstance(file, str):
+            file = Path(file)
+        if not file.is_file():
+            raise RuntimeError(f"xlsx file not found {file}")
+        self._file = file
+        self._book = openpyxl.load_workbook(file, read_only=read_only)
+
+    def keys(self) -> list[str]:
+        return [page.title for page in self._book.worksheets]
+
+    def stream_copy(self, into: str | Path) -> gen[tuple[str, int, list[str]]]:
+        if not isinstance(into, Path):
+            into = Path(into)
+
+        if not into.name.endswith(".xlsx"):
+            into = into / self._file.name
+
+        shutil.copy2(self._file, into)
+        the_copy = openpyxl.load_workbook(into)
+
+        for name in the_copy.sheetnames:
+            page = the_copy[name]
+
+            for r in range(page.max_row):
+                data = [page.cell(r + 1, c + 1).value for c in range(page.max_column)]
+                yield name, r, data
+                for c, val in enumerate(data):
+                    if val != page.cell(r + 1, c + 1).value:
+                        page.cell(r + 1, c + 1).value = val
+        the_copy.save(into)
+
+    @contextlib.contextmanager
+    def __getitem__(self, name: str) -> gen[Page]:
+
+        if len(name) > 31:
+            raise KeyError("excel limits you to 31 character page names")
+
+        class XLSXPage(Page):
+            def stream_rows(self) -> gen[tuple[int, list[str]]]:
+                for row in range(self._page.max_row):
+                    yield (
+                        row,
+                        [
+                            self._page.cell(row + 1, col + 1).value
+                            for col in range(self._page.max_column)
+                        ],
+                    )
+
+        yield XLSXPage(name, self._book[name])