rustfava 0.1.0__py3-none-any.whl

rustfava/core/tree.py

@@ -0,0 +1,265 @@
+"""Account balance trees."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from dataclasses import dataclass
+from operator import attrgetter
+from typing import TYPE_CHECKING
+
+from rustfava.beans.abc import Open
+from rustfava.beans.account import parent as account_parent
+from rustfava.context import g
+from rustfava.core.conversion import AT_COST
+from rustfava.core.conversion import AT_VALUE
+from rustfava.core.inventory import CounterInventory
+
+if TYPE_CHECKING:  # pragma: no cover
+    import datetime
+    from collections.abc import Iterable
+    from collections.abc import Sequence
+
+    from beancount.core import data
+
+    from rustfava.beans.abc import Directive
+    from rustfava.beans.prices import RustfavaPriceMap
+    from rustfava.beans.types import BeancountOptions
+    from rustfava.core.conversion import Conversion
+    from rustfava.core.inventory import SimpleCounterInventory
+
+
+@dataclass(frozen=True)
+class SerialisedTreeNode:
+    """A serialised TreeNode."""
+
+    account: str
+    balance: SimpleCounterInventory
+    balance_children: SimpleCounterInventory
+    children: Sequence[SerialisedTreeNode]
+    has_txns: bool
+    cost: SimpleCounterInventory | None = None
+    cost_children: SimpleCounterInventory | None = None
+
+
+class TreeNode:
+    """A node in the account tree."""
+
+    __slots__ = ("balance", "balance_children", "children", "has_txns", "name")
+
+    def __init__(self, name: str) -> None:
+        #: Account name.
+        self.name: str = name
+        #: A list of :class:`.TreeNode`, its children.
+        self.children: list[TreeNode] = []
+        #: The cumulative account balance.
+        self.balance_children = CounterInventory()
+        #: The account balance.
+        self.balance = CounterInventory()
+        #: Whether the account has any transactions.
+        self.has_txns = False
+
+    def serialise(
+        self,
+        conversion: Conversion,
+        prices: RustfavaPriceMap,
+        end: datetime.date | None,
+        *,
+        with_cost: bool = False,
+    ) -> SerialisedTreeNode:
+        """Serialise the account.
+
+        Args:
+            conversion: The conversion to use.
+            prices: The price map to use.
+            end: A date to use for cost conversions.
+            with_cost: Additionally convert to cost.
+        """
+        children = [
+            child.serialise(conversion, prices, end, with_cost=with_cost)
+            for child in sorted(self.children, key=attrgetter("name"))
+        ]
+        return (
+            SerialisedTreeNode(
+                self.name,
+                conversion.apply(self.balance, prices, end),
+                conversion.apply(self.balance_children, prices, end),
+                children,
+                self.has_txns,
+                AT_COST.apply(self.balance),
+                AT_COST.apply(self.balance_children),
+            )
+            if with_cost
+            else SerialisedTreeNode(
+                self.name,
+                conversion.apply(self.balance, prices, end),
+                conversion.apply(self.balance_children, prices, end),
+                children,
+                self.has_txns,
+            )
+        )
+
+    def serialise_with_context(self) -> SerialisedTreeNode:
+        """Serialise, getting all parameters from Flask context."""
+        return self.serialise(
+            g.conv,
+            g.ledger.prices,
+            g.filtered.end_date,
+            with_cost=g.conv == AT_VALUE,
+        )
+
+
+class Tree(dict[str, TreeNode]):
+    """Account tree.
+
+    Args:
+        entries: A list of entries to compute balances from.
+        create_accounts: A list of accounts that the tree should contain.
+    """
+
+    def __init__(
+        self,
+        entries: Iterable[Directive | data.Directive] | None = None,
+        create_accounts: list[str] | None = None,
+    ) -> None:
+        super().__init__(self)
+        self.get("", insert=True)
+        if create_accounts:
+            for account in create_accounts:
+                self.get(account, insert=True)
+        if entries:
+            account_balances: dict[str, CounterInventory]
+            account_balances = defaultdict(CounterInventory)
+            for entry in entries:
+                if isinstance(entry, Open):
+                    self.get(entry.account, insert=True)
+                for posting in getattr(entry, "postings", []):
+                    account_balances[posting.account].add_position(posting)
+
+            for name, balance in sorted(account_balances.items()):
+                self.insert(name, balance)
+
+    @property
+    def accounts(self) -> list[str]:
+        """The accounts in this tree."""
+        return sorted(self.keys())
+
+    def ancestors(self, name: str) -> Iterable[TreeNode]:
+        """Ancestors of an account.
+
+        Args:
+            name: An account name.
+
+        Yields:
+            The ancestors of the given account from the bottom up.
+        """
+        while name:
+            name = account_parent(name) or ""
+            yield self.get(name)
+
+    def insert(self, name: str, balance: CounterInventory) -> None:
+        """Insert account with a balance.
+
+        Insert account and update its balance and the balances of its
+        ancestors.
+
+        Args:
+            name: An account name.
+            balance: The balance of the account.
+        """
+        node = self.get(name, insert=True)
+        node.balance.add_inventory(balance)
+        node.balance_children.add_inventory(balance)
+        node.has_txns = True
+        for parent_node in self.ancestors(name):
+            parent_node.balance_children.add_inventory(balance)
+
+    def get(  # type: ignore[override]
+        self,
+        name: str,
+        *,
+        insert: bool = False,
+    ) -> TreeNode:
+        """Get an account.
+
+        Args:
+            name: An account name.
+            insert: If True, insert the name into the tree if it does not
+                exist.
+
+        Returns:
+            TreeNode: The account of that name or an empty account if the
+            account is not in the tree.
+        """
+        try:
+            return self[name]
+        except KeyError:
+            node = TreeNode(name)
+            if insert:
+                if name:
+                    parent = self.get(account_parent(name) or "", insert=True)
+                    parent.children.append(node)
+                self[name] = node
+            return node
+
+    def net_profit(
+        self,
+        options: BeancountOptions,
+        account_name: str,
+    ) -> TreeNode:
+        """Calculate the net profit.
+
+        Args:
+            options: The Beancount options.
+            account_name: The name to use for the account containing the net
+                profit.
+        """
+        income = self.get(options["name_income"])
+        expenses = self.get(options["name_expenses"])
+
+        net_profit = Tree()
+        net_profit.insert(
+            account_name,
+            income.balance_children + expenses.balance_children,
+        )
+
+        return net_profit.get(account_name)
+
+    def cap(self, options: BeancountOptions, unrealized_account: str) -> None:
+        """Transfer Income and Expenses, add conversions and unrealized gains.
+
+        Args:
+            options: The Beancount options.
+            unrealized_account: The name of the account to post unrealized
+                gains to (as a subaccount of Equity).
+        """
+        equity = options["name_equity"]
+        conversions = CounterInventory(
+            {
+                (currency, None): -number
+                for currency, number in AT_COST.apply(
+                    self.get("").balance_children
+                ).items()
+            },
+        )
+
+        # Add conversions
+        self.insert(
+            equity + ":" + options["account_current_conversions"],
+            conversions,
+        )
+
+        # Insert unrealized gains.
+        self.insert(
+            equity + ":" + unrealized_account,
+            -self.get("").balance_children,
+        )
+
+        # Transfer Income and Expenses
+        self.insert(
+            equity + ":" + options["account_current_earnings"],
+            self.get(options["name_income"]).balance_children,
+        )
+        self.insert(
+            equity + ":" + options["account_current_earnings"],
+            self.get(options["name_expenses"]).balance_children,
+        )

rustfava/core/watcher.py

@@ -0,0 +1,219 @@
+"""A simple file and folder watcher."""
+
+from __future__ import annotations
+
+import abc
+import atexit
+import logging
+import threading
+from os import walk
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from watchfiles import Change
+from watchfiles import DefaultFilter
+from watchfiles import watch
+
+if TYPE_CHECKING:  # pragma: no cover
+    import types
+    from collections.abc import Callable
+    from collections.abc import Iterable
+    from collections.abc import Sequence
+
+
+log = logging.getLogger(__name__)
+
+
+class _WatchfilesThread(threading.Thread):
+    """Class for the watchfiles watcher threads.
+
+    We use two separated threads since we want to recursively watch directories
+    and for paths, we need to watch the parent directory (to check changes done
+    by file replacements by some editors) non-recursively (for performance).
+    """
+
+    def __init__(
+        self,
+        paths: set[Path],
+        mtime: int,
+        *,
+        is_relevant: Callable[[Change, str], bool] | None = None,
+        recursive: bool = False,
+    ) -> None:
+        super().__init__(daemon=True)
+        self.paths = paths
+        self.mtime = mtime
+        self._is_relevant = is_relevant or DefaultFilter()
+        self._recursive = recursive
+        self._stop_event = threading.Event()
+
+    def stop(self) -> None:
+        """Set the stop event for watchfiles and join the thread."""
+        self._stop_event.set()
+        self.join()
+
+    def run(self) -> None:
+        """Watch for changes."""
+        atexit.register(self.stop)
+
+        for changes in watch(
+            *self.paths,
+            recursive=self._recursive,
+            stop_event=self._stop_event,
+            ignore_permission_denied=True,
+            watch_filter=self._is_relevant,
+        ):
+            for change_type, path_str in changes:
+                path = Path(path_str)
+                # move up the tree to an existing path
+                while not path.exists():
+                    path = path.parent
+                change_mtime = path.stat().st_mtime_ns
+                if change_type is Change.added:
+                    # check parent to get possibly newer timestamp of addition
+                    change_mtime = max(
+                        change_mtime, Path(path_str).parent.stat().st_mtime_ns
+                    )
+                self.mtime = max(change_mtime, self.mtime)
+            log.debug("new mtime: %s", self.mtime)
+
+
+class _FilesWatchfilesThread(_WatchfilesThread):
+    def __init__(self, files: set[Path], mtime: int) -> None:
+        paths = {f.parent for f in files}
+
+        def is_relevant(_c: Change, path: str) -> bool:
+            return Path(path) in files
+
+        super().__init__(
+            paths, mtime, is_relevant=is_relevant, recursive=False
+        )
+
+
+class WatcherBase(abc.ABC):
+    """ABC for rustfava ledger file watchers."""
+
+    last_checked: int
+    """Timestamp of the latest change noticed by the file watcher."""
+
+    last_notified: int
+    """Timestamp of the latest change that the watcher was notified of."""
+
+    @abc.abstractmethod
+    def update(self, files: Iterable[Path], folders: Iterable[Path]) -> None:
+        """Update the folders/files to watch.
+
+        Args:
+            files: A list of file paths.
+            folders: A list of paths to folders.
+        """
+
+    def check(self) -> bool:
+        """Check for changes.
+
+        Returns:
+            `True` if there was a file change in one of the files or folders,
+            `False` otherwise.
+        """
+        latest_mtime = max(self._get_latest_mtime(), self.last_notified)
+        has_higher_mtime = latest_mtime > self.last_checked
+        if has_higher_mtime:
+            self.last_checked = latest_mtime
+        return has_higher_mtime
+
+    def notify(self, path: Path) -> None:
+        """Notify the watcher of a change to a path."""
+        try:
+            change_mtime = Path(path).stat().st_mtime_ns
+        except FileNotFoundError:
+            change_mtime = max(self.last_notified, self.last_checked) + 1
+        self.last_notified = max(self.last_notified, change_mtime)
+
+    @abc.abstractmethod
+    def _get_latest_mtime(self) -> int:
+        """Get the latest change mtime."""
+
+
+class WatchfilesWatcher(WatcherBase):
+    """A file and folder watcher using the watchfiles library."""
+
+    def __init__(self) -> None:
+        self.last_checked = 0
+        self.last_notified = 0
+        self._paths: tuple[set[Path], set[Path]] | None = None
+        self._watchers: tuple[_WatchfilesThread, _WatchfilesThread] | None = (
+            None
+        )
+
+    def update(self, files: Iterable[Path], folders: Iterable[Path]) -> None:
+        """Update the folders/files to watch."""
+        files_set = {p.absolute() for p in files if p.exists()}
+        folders_set = {p.absolute() for p in folders if p.is_dir()}
+        new_paths = (files_set, folders_set)
+        if self._watchers and new_paths == self._paths:
+            self.check()
+            return
+        self._paths = new_paths
+        if self._watchers:
+            self._watchers[0].stop()
+            self._watchers[1].stop()
+        self._watchers = (
+            _FilesWatchfilesThread(files_set, self.last_checked),
+            _WatchfilesThread(folders_set, self.last_checked, recursive=True),
+        )
+        self._watchers[0].start()
+        self._watchers[1].start()
+        self.check()
+
+    def __enter__(self) -> None:
+        pass
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: types.TracebackType | None,
+    ) -> None:
+        if self._watchers:
+            self._watchers[0].stop()
+            self._watchers[1].stop()
+
+    def _get_latest_mtime(self) -> int:
+        return (
+            max(self._watchers[0].mtime, self._watchers[1].mtime)
+            if self._watchers
+            else 0
+        )
+
+
+class Watcher(WatcherBase):
+    """A simple file and folder watcher.
+
+    For folders, only checks mtime of the folder and all subdirectories.
+    So a file change won't be noticed, but only new/deleted files.
+    """
+
+    def __init__(self) -> None:
+        self.last_checked = 0
+        self.last_notified = 0
+        self._files: Sequence[Path] = []
+        self._folders: Sequence[Path] = []
+
+    def update(self, files: Iterable[Path], folders: Iterable[Path]) -> None:
+        """Update the folders/files to watch."""
+        self._files = list(files)
+        self._folders = list(folders)
+        self.check()
+
+    def _mtimes(self) -> Iterable[int]:
+        for path in self._files:
+            try:
+                yield path.stat().st_mtime_ns
+            except FileNotFoundError:
+                yield max(self.last_notified, self.last_checked) + 1
+        for path in self._folders:
+            for dirpath, _, _ in walk(path):
+                yield Path(dirpath).stat().st_mtime_ns
+
+    def _get_latest_mtime(self) -> int:
+        return max(self._mtimes())

rustfava/ext/__init__.py

@@ -0,0 +1,232 @@
+"""rustfava's extension system."""
+
+from __future__ import annotations
+
+import ast
+import importlib
+import inspect
+import sys
+from functools import cached_property
+from pathlib import Path
+from typing import Any
+from typing import TYPE_CHECKING
+
+import jinja2
+from flask import current_app
+
+from rustfava.helpers import BeancountError
+
+if TYPE_CHECKING:  # pragma: no cover
+    from collections.abc import Callable
+    from typing import TypeVar
+
+    from flask.wrappers import Response
+
+    from rustfava.beans.abc import Directive
+    from rustfava.core import RustfavaLedger
+
+
+class RustfavaExtensionError(BeancountError):
+    """Error in one of Fava's extensions."""
+
+
+class JinjaLoaderMissingError(ValueError):  # noqa: D101
+    def __init__(self) -> None:  # pragma: no cover
+        super().__init__("Expected Flask app to have jinja_loader.")
+
+
+class ExtensionConfigError(ValueError):  # noqa: D101
+    def __init__(self, error: SyntaxError, config: str) -> None:
+        super().__init__(
+            f"Could not load extension config: {error} in '{config}'."
+        )
+
+
+class RustfavaExtensionBase:
+    """Base class for extensions for Fava.
+
+    Any extension should inherit from this class. :func:`find_extension` will
+    discover all subclasses of this class in the specified modules.
+    """
+
+    #: Name for a HTML report for this extension.
+    report_title: str | None = None
+
+    #: Whether this extension includes a Javascript module.
+    has_js_module: bool = False
+
+    config: Any
+
+    endpoints: dict[tuple[str, str], Callable[[RustfavaExtensionBase], Any]]
+
+    def __init__(self, ledger: RustfavaLedger, config: str | None = None) -> None:
+        """Initialise extension.
+
+        Args:
+            ledger: Input ledger file.
+            config: Configuration options string passed from the
+                    beancount file's 'fava-extension' line.
+
+        Raises:
+            ExtensionConfigError: If the config cannot be parsed.
+        """
+        self.endpoints = {}
+        self.config = None
+
+        # Go through each of the subclass's functions to find the ones
+        # marked as endpoints by @extension_endpoint
+        for _, func in inspect.getmembers(self.__class__, inspect.isfunction):
+            if hasattr(func, "endpoint_key"):
+                name, methods = func.endpoint_key
+                for method in methods:
+                    self.endpoints[name, method] = func
+
+        self.ledger = ledger
+        if config:
+            try:
+                self.config = ast.literal_eval(config)
+            except SyntaxError as error:
+                raise ExtensionConfigError(error, config) from error
+
+    @property
+    def name(self) -> str:
+        """Unique name of this extension."""
+        return self.__class__.__qualname__
+
+    @property
+    def extension_dir(self) -> Path:
+        """Directory to look for templates directory and Javascript code."""
+        return Path(inspect.getfile(self.__class__)).parent
+
+    @cached_property
+    def jinja_env(self) -> jinja2.Environment:
+        """Jinja env for this extension."""
+        if not current_app.jinja_loader:  # pragma: no cover
+            raise JinjaLoaderMissingError
+        ext_loader = jinja2.FileSystemLoader(self.extension_dir / "templates")
+        loader = jinja2.ChoiceLoader([ext_loader, current_app.jinja_loader])
+        return current_app.jinja_env.overlay(loader=loader)
+
+    def after_load_file(self) -> None:
+        """Run after a ledger file has been loaded."""
+
+    def before_request(self) -> None:
+        """Run before each client request."""
+
+    def after_entry_modified(self, entry: Directive, new_lines: str) -> None:
+        """Run after an `entry` has been modified."""
+
+    def after_insert_entry(self, entry: Directive) -> None:
+        """Run after an `entry` has been inserted."""
+
+    def after_delete_entry(self, entry: Directive) -> None:
+        """Run after an `entry` has been deleted."""
+
+    def after_insert_metadata(
+        self,
+        entry: Directive,
+        key: str,
+        value: str,
+    ) -> None:
+        """Run after metadata (key: value) was added to an entry."""
+
+    def after_write_source(self, path: str, source: str) -> None:
+        """Run after `source` has been written to path."""
+
+
+def find_extensions(
+    base_path: Path,
+    name: str,
+) -> tuple[list[type[RustfavaExtensionBase]], list[RustfavaExtensionError]]:
+    """Find extensions in a module.
+
+    Args:
+        base_path: The module can be relative to this path.
+        name: The name of the module containing the extensions.
+
+    Returns:
+        A tuple (classes, errors) where classes is a list of subclasses of
+        :class:`RustfavaExtensionBase` found in ``name``.
+    """
+    classes = []
+
+    sys.path.insert(0, str(base_path))
+    try:
+        module = importlib.import_module(name)
+    except ImportError as err:
+        error = RustfavaExtensionError(
+            None,
+            f'Importing module "{name}" failed.\nError: "{err.msg}"',
+            None,
+        )
+        return (
+            [],
+            [error],
+        )
+    for _, obj in inspect.getmembers(module, inspect.isclass):
+        if issubclass(obj, RustfavaExtensionBase) and obj != RustfavaExtensionBase:
+            classes.append(obj)
+    sys.path.pop(0)
+
+    if not classes:
+        error = RustfavaExtensionError(
+            None,
+            f'Module "{name}" contains no extensions.',
+            None,
+        )
+        return (
+            [],
+            [error],
+        )
+
+    return classes, []
+
+
+if TYPE_CHECKING:  # pragma: no cover
+    T = TypeVar("T", bound=RustfavaExtensionBase)
+
+
+def extension_endpoint(
+    func_or_endpoint_name: (Callable[[T], Any] | str | None) = None,
+    methods: list[str] | None = None,
+) -> (
+    Callable[[T], Response]
+    | Callable[
+        [Callable[[T], Response]],
+        Callable[[T], Response],
+    ]
+):
+    """Decorator to mark a function as an endpoint.
+
+    Can be used as `@extension_endpoint` or
+    `@extension_endpoint(endpoint_name, methods)`.
+
+    When used as @extension_endpoint, the endpoint name is the name of the
+    function and methods is "GET".
+
+    When used as @extension_endpoint(endpoint_name, methods), the given
+    endpoint name and methods are used, but both are optional. If
+    endpoint_name is None, default to the function name, and if methods
+    is None, default to "GET".
+    """
+    endpoint_name = (
+        func_or_endpoint_name
+        if isinstance(func_or_endpoint_name, str)
+        else None
+    )
+
+    def decorator(
+        func: Callable[[T], Response],
+    ) -> Callable[[T], Response]:
+        f: Any = func
+        f.endpoint_key = (  # ty:ignore[unresolved-attribute]
+            endpoint_name or func.__name__,  # ty:ignore[unresolved-attribute]
+            methods or ["GET"],
+        )
+        return func
+
+    return (
+        decorator(func_or_endpoint_name)  # ty:ignore[invalid-argument-type]
+        if callable(func_or_endpoint_name)
+        else decorator
+    )