rustfava 0.1.0__py3-none-any.whl

rustfava/core/ingest.py

@@ -0,0 +1,509 @@
+"""Ingest helper functions."""
+
+from __future__ import annotations
+
+import datetime
+import os
+import traceback
+from dataclasses import dataclass
+from functools import wraps
+from inspect import get_annotations
+from inspect import signature
+from os import altsep
+from os import sep
+from pathlib import Path
+from runpy import run_path
+from typing import TYPE_CHECKING
+
+from beangulp import Importer
+
+try:  # pragma: no cover
+    from beancount.ingest import cache  # type: ignore[import-not-found]
+    from beancount.ingest import extract
+
+    DEFAULT_HOOKS = [extract.find_duplicate_entries]
+except ImportError:
+    from beangulp import cache
+
+    DEFAULT_HOOKS = []
+
+from rustfava.beans.ingest import BeanImporterProtocol
+from rustfava.core.file import _incomplete_sortkey
+from rustfava.core.module_base import FavaModule
+from rustfava.helpers import BeancountError
+from rustfava.helpers import RustfavaAPIError
+from rustfava.util.date import local_today
+
+if TYPE_CHECKING:  # pragma: no cover
+    from collections.abc import Callable
+    from collections.abc import Iterable
+    from collections.abc import Mapping
+    from collections.abc import Sequence
+    from typing import Any
+    from typing import ParamSpec
+    from typing import TypeVar
+
+    from rustfava.beans.abc import Directive
+    from rustfava.beans.ingest import FileMemo
+    from rustfava.core import RustfavaLedger
+
+    HookOutput = (
+        list[tuple[str, list[Directive], str, BeanImporterProtocol | Importer]]
+        | list[tuple[str, list[Directive]]]
+    )
+    Hooks = Sequence[Callable[[HookOutput, Sequence[Directive]], HookOutput]]
+
+    P = ParamSpec("P")
+    T = TypeVar("T")
+
+
+class IngestError(BeancountError):
+    """An error with one of the importers."""
+
+
+class ImporterMethodCallError(RustfavaAPIError):
+    """Error calling one of the importer methods."""
+
+    def __init__(self) -> None:
+        super().__init__(
+            f"Error calling method on importer:\n\n{traceback.format_exc()}"
+        )
+
+
+class ImporterInvalidTypeError(RustfavaAPIError):
+    """One of the importer methods returned an unexpected type."""
+
+    def __init__(self, attr: str, expected: type[Any], actual: Any) -> None:
+        super().__init__(
+            f"Got unexpected type from importer as {attr}:"
+            f" expected {expected!s}, got {type(actual)!s}:"
+        )
+
+
+class ImporterExtractError(ImporterMethodCallError):
+    """Error calling extract for importer."""
+
+
+class MissingImporterConfigError(RustfavaAPIError):
+    """Missing import-config option."""
+
+    def __init__(self) -> None:
+        super().__init__("Missing import-config option")
+
+
+class MissingImporterDirsError(RustfavaAPIError):
+    """You need to set at least one imports-dir."""
+
+    def __init__(self) -> None:
+        super().__init__("You need to set at least one imports-dir.")
+
+
+class ImportConfigLoadError(RustfavaAPIError):
+    """Error on loading the import config."""
+
+
+IGNORE_DIRS = {
+    ".cache",
+    ".git",
+    ".hg",
+    ".idea",
+    ".svn",
+    ".tox",
+    ".venv",
+    "__pycache__",
+    "node_modules",
+}
+
+
+def walk_dir(directory: Path) -> Iterable[Path]:
+    """Walk through all files in dir.
+
+    Ignores common dot-directories like .git, .cache. .venv, see IGNORE_DIRS.
+
+    Args:
+        directory: The directory to start in.
+
+    Yields:
+        All full paths under directory, ignoring some directories.
+    """
+    for root, dirs, filenames in os.walk(directory):
+        dirs[:] = sorted(d for d in dirs if d not in IGNORE_DIRS)
+        root_path = Path(root)
+        for filename in sorted(filenames):
+            yield root_path / filename
+
+
+# Keep our own cache to also keep track of file mtimes
+_CACHE: dict[Path, tuple[int, FileMemo]] = {}
+
+
+def get_cached_file(path: Path) -> FileMemo:
+    """Get a cached FileMemo.
+
+    This checks the file's mtime before getting it from the Cache.
+    In addition to using the beangulp cache.
+    """
+    mtime = path.stat().st_mtime_ns
+    filename = str(path)
+    cached = _CACHE.get(path)
+    if cached:
+        mtime_cached, memo_cached = cached
+        if mtime <= mtime_cached:  # pragma: no cover
+            return memo_cached
+    memo: FileMemo = cache._FileMemo(filename)  # noqa: SLF001
+    cache._CACHE[filename] = memo  # noqa: SLF001
+    _CACHE[path] = (mtime, memo)
+    return memo
+
+
+@dataclass(frozen=True)
+class FileImportInfo:
+    """Info about one file/importer combination."""
+
+    importer_name: str
+    account: str
+    date: datetime.date
+    name: str
+
+
+@dataclass(frozen=True)
+class FileImporters:
+    """Importers for a file."""
+
+    name: str
+    basename: str
+    importers: list[FileImportInfo]
+
+
+def _catch_any(func: Callable[P, T]) -> Callable[P, T]:
+    """Helper to catch any exception that might be raised by the importer."""
+
+    @wraps(func)
+    def wrapper(*args: P.args, **kwds: P.kwargs) -> T:
+        try:
+            return func(*args, **kwds)
+        except Exception as err:
+            if isinstance(err, ImporterInvalidTypeError):
+                raise
+            raise ImporterMethodCallError from err
+
+    return wrapper
+
+
+def _assert_type(attr: str, value: T, type_: type[T]) -> T:
+    """Helper to validate types return by importer methods."""
+    if not isinstance(value, type_):
+        raise ImporterInvalidTypeError(attr, type_, value)
+    return value
+
+
+class WrappedImporter:
+    """A wrapper to safely call importer methods."""
+
+    importer: BeanImporterProtocol | Importer
+
+    def __init__(self, importer: BeanImporterProtocol | Importer) -> None:
+        self.importer = importer
+
+    @property
+    @_catch_any
+    def name(self) -> str:
+        """Get the name of the importer."""
+        importer = self.importer
+        name = (
+            importer.name
+            if isinstance(importer, Importer)
+            else importer.name()
+        )
+        return _assert_type("name", name, str)
+
+    @_catch_any
+    def identify(self: WrappedImporter, path: Path) -> bool:
+        """Whether the importer is matching the file."""
+        importer = self.importer
+        matches = (
+            importer.identify(str(path))
+            if isinstance(importer, Importer)
+            else importer.identify(get_cached_file(path))
+        )
+        return _assert_type("identify", matches, bool)
+
+    @_catch_any
+    def file_import_info(self, path: Path) -> FileImportInfo:
+        """Generate info about a file with an importer."""
+        importer = self.importer
+        if isinstance(importer, Importer):
+            str_path = str(path)
+            account = importer.account(str_path)
+            date = importer.date(str_path)
+            filename = importer.filename(str_path)
+        else:
+            file = get_cached_file(path)
+            account = importer.file_account(file)
+            date = importer.file_date(file)
+            filename = importer.file_name(file)
+
+        return FileImportInfo(
+            self.name,
+            _assert_type("account", account or "", str),
+            _assert_type("date", date or local_today(), datetime.date),
+            _assert_type("filename", filename or path.name, str),
+        )
+
+
+# Copied here from beangulp to minimise the imports.
+_FILE_TOO_LARGE_THRESHOLD = 8 * 1024 * 1024
+
+
+def find_imports(
+    config: Sequence[WrappedImporter], directory: Path
+) -> Iterable[FileImporters]:
+    """Pair files and matching importers.
+
+    Yields:
+        For each file in directory, a pair of its filename and the matching
+        importers.
+    """
+    for path in walk_dir(directory):
+        stat = path.stat()
+        if stat.st_size > _FILE_TOO_LARGE_THRESHOLD:  # pragma: no cover
+            continue
+
+        importers = [
+            importer.file_import_info(path)
+            for importer in config
+            if importer.identify(path)
+        ]
+        yield FileImporters(
+            name=str(path), basename=path.name, importers=importers
+        )
+
+
+def extract_from_file(
+    wrapped_importer: WrappedImporter,
+    path: Path,
+    existing_entries: Sequence[Directive],
+) -> list[Directive]:
+    """Import entries from a document.
+
+    Args:
+      wrapped_importer: The importer instance to handle the document.
+      path: Filesystem path to the document.
+      existing_entries: Existing entries.
+
+    Returns:
+      The list of imported entries.
+    """
+    filename = str(path)
+    importer = wrapped_importer.importer
+    if isinstance(importer, Importer):
+        entries = importer.extract(filename, existing=existing_entries)
+    else:
+        file = get_cached_file(path)
+        entries = (
+            importer.extract(file, existing_entries=existing_entries)
+            if "existing_entries" in signature(importer.extract).parameters
+            else importer.extract(file)
+        ) or []
+
+    if hasattr(importer, "sort"):
+        importer.sort(entries)
+    else:
+        entries.sort(key=_incomplete_sortkey)
+    if isinstance(importer, Importer):
+        importer.deduplicate(entries, existing=existing_entries)
+    return entries
+
+
+def load_import_config(
+    module_path: Path,
+) -> tuple[Mapping[str, WrappedImporter], Hooks]:
+    """Load the given import config and extract importers and hooks.
+
+    Args:
+        module_path: Path to the import config.
+
+    Returns:
+        A pair of the importers (by name) and the list of hooks.
+    """
+    try:
+        mod = run_path(str(module_path))
+    except Exception as error:  # pragma: no cover
+        message = traceback.format_exc()
+        raise ImportConfigLoadError(message) from error
+
+    if "CONFIG" not in mod:
+        msg = "CONFIG is missing"
+        raise ImportConfigLoadError(msg)
+    if not isinstance(mod["CONFIG"], list):  # pragma: no cover
+        msg = "CONFIG is not a list"
+        raise ImportConfigLoadError(msg)
+
+    config = mod["CONFIG"]
+    hooks = DEFAULT_HOOKS
+    if "HOOKS" in mod:  # pragma: no cover
+        hooks = mod["HOOKS"]
+        if not isinstance(hooks, list) or not all(
+            callable(fn) for fn in hooks
+        ):
+            msg = "HOOKS is not a list of callables"
+            raise ImportConfigLoadError(msg)
+    importers = {}
+    for importer in config:
+        if not isinstance(
+            importer, (BeanImporterProtocol, Importer)
+        ):  # pragma: no cover
+            name = importer.__class__.__name__
+            msg = (
+                f"Importer class '{name}' in '{module_path}' does "
+                "not satisfy importer protocol"
+            )
+            raise ImportConfigLoadError(msg)
+        wrapped_importer = WrappedImporter(importer)
+        if wrapped_importer.name in importers:
+            msg = f"Duplicate importer name found: {wrapped_importer.name}"
+            raise ImportConfigLoadError(msg)
+        importers[wrapped_importer.name] = wrapped_importer
+    return importers, hooks
+
+
+class IngestModule(FavaModule):
+    """Exposes ingest functionality."""
+
+    def __init__(self, ledger: RustfavaLedger) -> None:
+        super().__init__(ledger)
+        self.importers: Mapping[str, WrappedImporter] = {}
+        self.hooks: Hooks = []
+        self.mtime: int | None = None
+        self.errors: list[IngestError] = []
+
+    @property
+    def module_path(self) -> Path | None:
+        """The path to the importer configuration."""
+        config_path = self.ledger.fava_options.import_config
+        if not config_path:
+            return None
+        return self.ledger.join_path(config_path)
+
+    def _error(self, msg: str) -> None:
+        self.errors.append(
+            IngestError(
+                {"filename": str(self.module_path), "lineno": 0},
+                msg,
+                None,
+            ),
+        )
+
+    def load_file(self) -> None:  # noqa: D102
+        self.errors = []
+        module_path = self.module_path
+        if module_path is None:
+            return
+
+        if not module_path.exists():
+            self._error("Import config does not exist")
+            return
+
+        new_mtime = module_path.stat().st_mtime_ns
+        if new_mtime == self.mtime:
+            return
+
+        try:
+            self.importers, self.hooks = load_import_config(module_path)
+            self.mtime = new_mtime
+        except RustfavaAPIError as error:  # pragma: no cover
+            msg = f"Error in import config '{module_path}': {error!s}"
+            self._error(msg)
+
+    def import_data(self) -> list[FileImporters]:
+        """Identify files and importers that can be imported.
+
+        Returns:
+            A list of :class:`.FileImportInfo`.
+        """
+        if not self.importers:
+            return []
+
+        importers = list(self.importers.values())
+
+        ret: list[FileImporters] = []
+        for directory in self.ledger.fava_options.import_dirs:
+            full_path = self.ledger.join_path(directory)
+            ret.extend(find_imports(importers, full_path))
+
+        return ret
+
+    def extract(self, filename: str, importer_name: str) -> list[Directive]:
+        """Extract entries from filename with the specified importer.
+
+        Args:
+            filename: The full path to a file.
+            importer_name: The name of an importer that matched the file.
+
+        Returns:
+            A list of new imported entries.
+        """
+        if not self.module_path:
+            raise MissingImporterConfigError
+
+        # reload (if changed)
+        self.load_file()
+
+        try:
+            path = Path(filename)
+            importer = self.importers[importer_name]
+            new_entries = extract_from_file(
+                importer,
+                path,
+                existing_entries=self.ledger.all_entries,
+            )
+        except Exception as exc:
+            raise ImporterExtractError from exc
+
+        for hook_fn in self.hooks:
+            annotations = get_annotations(hook_fn)
+            if any("Importer" in a for a in annotations.values()):
+                importer_info = importer.file_import_info(path)
+                new_entries_list: HookOutput = [
+                    (
+                        filename,
+                        new_entries,
+                        importer_info.account,
+                        importer.importer,
+                    )
+                ]
+            else:
+                new_entries_list = [(filename, new_entries)]
+
+            new_entries_list = hook_fn(
+                new_entries_list,
+                self.ledger.all_entries,
+            )
+
+            new_entries = new_entries_list[0][1]
+
+        return new_entries
+
+
+def filepath_in_primary_imports_folder(
+    filename: str,
+    ledger: RustfavaLedger,
+) -> Path:
+    """File path for a document to upload to the primary import folder.
+
+    Args:
+        filename: The filename of the document.
+        ledger: The RustfavaLedger.
+
+    Returns:
+        The path that the document should be saved at.
+    """
+    primary_imports_folder = next(iter(ledger.fava_options.import_dirs), None)
+    if primary_imports_folder is None:
+        raise MissingImporterDirsError
+
+    filename = filename.replace(sep, " ")
+    if altsep:  # pragma: no cover
+        filename = filename.replace(altsep, " ")
+
+    return ledger.join_path(primary_imports_folder, filename)

rustfava/core/inventory.py

@@ -0,0 +1,167 @@
+"""Alternative implementation of Beancount's Inventory."""
+
+from __future__ import annotations
+
+from decimal import Decimal
+from typing import NamedTuple
+from typing import TYPE_CHECKING
+
+from rustfava.beans.protocols import Cost
+from rustfava.beans.str import cost_to_string
+
+if TYPE_CHECKING:  # pragma: no cover
+    import datetime
+    from collections.abc import Callable
+    from collections.abc import Iterator
+    from typing import Concatenate
+    from typing import ParamSpec
+
+    from rustfava.beans.protocols import Amount
+    from rustfava.beans.protocols import Position
+
+    P = ParamSpec("P")
+
+
+ZERO = Decimal()
+InventoryKey = tuple[str, Cost | None]
+
+
+class _Amount(NamedTuple):
+    number: Decimal
+    currency: str
+
+
+class _Cost(NamedTuple):
+    number: Decimal
+    currency: str
+    date: datetime.date
+    label: str | None
+
+
+class _Position(NamedTuple):
+    units: Amount
+    cost: Cost | None
+
+
+class SimpleCounterInventory(dict[str, Decimal]):
+    """A simple inventory mapping just strings to numbers."""
+
+    def is_empty(self) -> bool:
+        """Check if the inventory is empty."""
+        return not bool(self)
+
+    def add(self, key: str, number: Decimal) -> None:
+        """Add a number to key."""
+        new_num = number + self.get(key, ZERO)
+        if new_num == ZERO:
+            self.pop(key, None)
+        else:
+            self[key] = new_num
+
+    def __iter__(self) -> Iterator[str]:
+        raise NotImplementedError
+
+    def __neg__(self) -> SimpleCounterInventory:
+        return SimpleCounterInventory({key: -num for key, num in self.items()})
+
+    def reduce(
+        self,
+        reducer: Callable[Concatenate[Position, P], Amount],
+        *args: P.args,
+        **_kwargs: P.kwargs,
+    ) -> SimpleCounterInventory:
+        """Reduce inventory."""
+        counter = SimpleCounterInventory()
+        for currency, number in self.items():
+            pos = _Position(_Amount(number, currency), None)
+            amount = reducer(pos, *args)  # type: ignore[call-arg]
+            counter.add(amount.currency, amount.number)
+        return counter
+
+
+class CounterInventory(dict[InventoryKey, Decimal]):
+    """A lightweight inventory.
+
+    This is intended as a faster alternative to Beancount's Inventory class.
+    Due to not using a list, for inventories with a lot of different positions,
+    inserting is much faster.
+
+    The keys should be tuples ``(currency, cost)``.
+    """
+
+    def is_empty(self) -> bool:
+        """Check if the inventory is empty."""
+        return not bool(self)
+
+    def add(self, key: InventoryKey, number: Decimal) -> None:
+        """Add a number to key."""
+        new_num = number + self.get(key, ZERO)
+        if new_num == ZERO:
+            self.pop(key, None)
+        else:
+            self[key] = new_num
+
+    def __iter__(self) -> Iterator[InventoryKey]:
+        raise NotImplementedError
+
+    def to_strings(self) -> list[str]:
+        """Print as a list of strings (e.g. for snapshot tests)."""
+        strings = []
+        for (currency, cost), number in self.items():
+            if cost is None:
+                strings.append(f"{number} {currency}")
+            else:
+                cost_str = cost_to_string(cost)
+                strings.append(f"{number} {currency} {{{cost_str}}}")
+        return strings
+
+    def reduce(
+        self,
+        reducer: Callable[Concatenate[Position, P], Amount],
+        *args: P.args,
+        **_kwargs: P.kwargs,
+    ) -> SimpleCounterInventory:
+        """Reduce inventory.
+
+        Note that this returns a simple :class:`CounterInventory` with just
+        currencies as keys.
+        """
+        counter = SimpleCounterInventory()
+        for (currency, cost), number in self.items():
+            pos = _Position(_Amount(number, currency), cost)
+            amount = reducer(pos, *args)  # type: ignore[call-arg]
+            counter.add(amount.currency, amount.number)
+        return counter
+
+    def add_amount(self, amount: Amount, cost: Cost | None = None) -> None:
+        """Add an Amount to the inventory."""
+        key = (amount.currency, cost)
+        self.add(key, amount.number)
+
+    def add_position(self, pos: Position) -> None:
+        """Add a Position or Posting to the inventory."""
+        # Skip positions with missing units (can happen with parse errors)
+        if pos.units is None:
+            return
+        self.add_amount(pos.units, pos.cost)
+
+    def __neg__(self) -> CounterInventory:
+        return CounterInventory({key: -num for key, num in self.items()})
+
+    def __add__(self, other: CounterInventory) -> CounterInventory:
+        counter = CounterInventory(self)
+        counter.add_inventory(other)
+        return counter
+
+    def add_inventory(self, counter: CounterInventory) -> None:
+        """Add another :class:`CounterInventory`."""
+        if not self:
+            self.update(counter)
+        else:
+            self_get = self.get
+            for key, num in counter.items():
+                new_num = num + self_get(key, ZERO)
+                if new_num == ZERO:
+                    self.pop(key, None)
+                else:
+                    self[key] = new_num