nomenklatura-mpt 4.1.9__py3-none-any.whl

nomenklatura/__init__.py

@@ -0,0 +1,11 @@
+from nomenklatura.resolver import Resolver
+from nomenklatura.store import Store, View
+from nomenklatura.index import Index
+
+__version__ = "4.1.9"
+__all__ = [
+    "Resolver",
+    "Index",
+    "Store",
+    "View",
+]

nomenklatura/cache.py

@@ -0,0 +1,194 @@
+import math
+import json
+import logging
+from random import randint
+from dataclasses import dataclass
+from typing import Any, cast, Dict, Optional, Union, Generator
+from datetime import datetime, timedelta
+from sqlalchemy import MetaData
+from sqlalchemy import Table, Column, DateTime, Unicode
+from sqlalchemy.engine import Engine, Connection, Transaction
+from sqlalchemy.future import select
+from sqlalchemy.sql.expression import delete
+from sqlalchemy.exc import OperationalError, InvalidRequestError
+from sqlalchemy.dialects.postgresql import insert as upsert
+from rigour.time import naive_now
+from followthemoney import Dataset
+
+from nomenklatura.db import get_engine, get_metadata
+
+
+log = logging.getLogger(__name__)
+Value = Union[str, None]
+
+
+@dataclass
+class CacheValue:
+    key: str
+    dataset: Optional[str]
+    text: Value
+    timestamp: datetime
+
+
+def randomize_cache(days: int) -> timedelta:
+    min_cache = max(1, math.ceil(days * 0.5))
+    max_cache = math.ceil(days * 1.3)
+    return timedelta(days=randint(min_cache, max_cache))
+
+
+class Cache(object):
+    def __init__(
+        self, engine: Engine, metadata: MetaData, dataset: Dataset, create: bool = False
+    ) -> None:
+        self.dataset = dataset
+        self._engine = engine
+        self._conn: Optional[Connection] = None
+        self._transaction: Optional[Transaction] = None
+        self._table = Table(
+            "cache",
+            metadata,
+            Column("key", Unicode(), primary_key=True),
+            Column("text", Unicode(), nullable=True),
+            Column("dataset", Unicode(), nullable=False),
+            Column("timestamp", DateTime, index=True),
+            extend_existing=True,
+        )
+        if create:
+            metadata.create_all(bind=engine, checkfirst=True, tables=[self._table])
+
+        self._preload: Dict[str, CacheValue] = {}
+
+    @property
+    def conn(self) -> Connection:
+        if self._conn is None:
+            self._conn = self._engine.connect()
+            self._transaction = self._conn.begin()
+        return self._conn
+
+    def set(self, key: str, value: Value) -> None:
+        self._preload.pop(key, None)
+        cache = {
+            "timestamp": naive_now(),
+            "key": key,
+            "dataset": self.dataset.name,
+            "text": value,
+        }
+        try:
+            istmt = upsert(self._table).values(cache)
+            values = dict(
+                timestamp=istmt.excluded.timestamp,
+                text=istmt.excluded.text,
+                dataset=istmt.excluded.dataset,
+            )
+            stmt = istmt.on_conflict_do_update(index_elements=["key"], set_=values)
+            self.conn.execute(stmt)
+        except (OperationalError, InvalidRequestError) as exc:
+            log.exception("Error while saving to cache: %s" % exc)
+            self.reset()
+
+    def set_json(self, key: str, value: Any) -> None:
+        return self.set(key, json.dumps(value))
+
+    def get(self, key: str, max_age: Optional[int] = None) -> Optional[Value]:
+        if max_age is not None and max_age < 1:
+            return None
+
+        cache_cutoff = None
+        if max_age is not None:
+            cache_cutoff = naive_now() - randomize_cache(max_age)
+
+        cache = self._preload.get(key)
+        if cache is not None:
+            if cache_cutoff is not None and cache.timestamp < cache_cutoff:
+                return None
+            return cache.text
+
+        q = select(self._table.c.text)
+        q = q.filter(self._table.c.key == key)
+        if cache_cutoff is not None:
+            q = q.filter(self._table.c.timestamp > cache_cutoff)
+        q = q.order_by(self._table.c.timestamp.desc())
+        q = q.limit(1)
+        try:
+            result = self.conn.execute(q)
+            row = result.fetchone()
+        except InvalidRequestError as ire:
+            log.exception("Cache fetch error: %s", ire)
+            self.reset()
+            return None
+        if row is not None:
+            return cast(Optional[str], row.text)
+        return None
+
+    def get_json(self, key: str, max_age: Optional[int] = None) -> Optional[Any]:
+        text = self.get(key, max_age=max_age)
+        if text is None:
+            return None
+        return json.loads(text)
+
+    def has(self, key: str) -> bool:
+        return self.get(key) is not None
+
+    def delete(self, key: str) -> None:
+        self._preload.pop(key, None)
+        pq = delete(self._table)
+        pq = pq.where(self._table.c.key == key)
+        try:
+            self.conn.execute(pq)
+        except InvalidRequestError as ire:
+            log.exception("Cache delete error: %s", ire)
+            self.reset()
+            return None
+
+    def all(self, like: Optional[str]) -> Generator[CacheValue, None, None]:
+        q = select(self._table)
+        if like is not None:
+            q = q.filter(self._table.c.key.like(like))
+
+        result = self.conn.execute(q)
+        for row in result.yield_per(10000):
+            yield CacheValue(row.key, row.dataset, row.text, row.timestamp)
+
+    def preload(self, like: Optional[str] = None) -> None:
+        log.info("Pre-loading cache: %r", like)
+        for cache in self.all(like=like):
+            self._preload[cache.key] = cache
+
+    def clear(self) -> None:
+        try:
+            pq = delete(self._table)
+            pq = pq.where(self._table.c.dataset == self.dataset.name)
+            self.conn.execute(pq)
+        except InvalidRequestError:
+            log.exception("Cannot clear cache from database")
+            self.reset()
+
+    def reset(self) -> None:
+        if self._conn is not None:
+            self._conn.close()
+        self._conn = None
+        self._transaction = None
+
+    def flush(self) -> None:
+        # log.info("Flushing cache.")
+        if self._transaction is not None:
+            try:
+                self._transaction.commit()
+            except InvalidRequestError:
+                log.exception("Transaction was failed, cannot store cache state.")
+        self.reset()
+
+    def close(self) -> None:
+        self.flush()
+
+    def __repr__(self) -> str:
+        return f"<Cache({self._table!r})>"
+
+    def __hash__(self) -> int:
+        return hash((self.dataset.name, self._table.name))
+
+    @classmethod
+    def make_default(cls, dataset: Dataset) -> "Cache":
+        engine = get_engine()
+        metadata = get_metadata()
+        return cls(engine, metadata, dataset, create=True)

nomenklatura/cli.py

@@ -0,0 +1,260 @@
+import os
+import shutil
+import yaml
+import click
+import logging
+from pathlib import Path
+from typing import Generator, Optional, Tuple
+from followthemoney import Dataset, ValueEntity, StatementEntity as Entity
+from followthemoney.statement import Statement, CSV, FORMATS
+from followthemoney.statement import write_statements, read_path_statements
+from followthemoney.cli.util import path_writer, InPath, OutPath
+from followthemoney.cli.util import path_entities, write_entity
+from followthemoney.cli.aggregate import sorted_aggregate
+
+from nomenklatura.cache import Cache
+from nomenklatura.matching import train_v1_matcher, train_erun_matcher
+from nomenklatura.store import load_entity_file_store
+from nomenklatura.resolver import Resolver, Linker
+from nomenklatura.enrich import Enricher, make_enricher, match, enrich
+from nomenklatura.matching import get_algorithm, DefaultAlgorithm
+from nomenklatura.xref import xref as run_xref
+from nomenklatura.tui import dedupe_ui
+from nomenklatura.matching.bench import bench_matcher
+
+INDEX_SEGMENT = "xref-index"
+
+log = logging.getLogger(__name__)
+
+ResPath = click.Path(dir_okay=False, writable=True, path_type=Path)
+
+
+def _load_enricher(path: Path) -> Tuple[Dataset, Enricher[Dataset]]:
+    with open(path, "r") as fh:
+        data = yaml.safe_load(fh)
+        dataset = Dataset.make(data)
+        cache = Cache.make_default(dataset)
+        enricher = make_enricher(dataset, cache, data)
+        if enricher is None:
+            raise TypeError("Could not load enricher")
+        return dataset, enricher
+
+
+def _get_linker() -> Linker[Entity]:
+    resolver = Resolver[Entity].make_default()
+    linker = resolver.get_linker()
+    resolver.close()
+    return linker
+
+
+@click.group(help="Nomenklatura data integration")
+def cli() -> None:
+    logging.basicConfig(level=logging.INFO)
+
+
+@cli.command("xref", help="Generate dedupe candidates")
+@click.argument("path", type=InPath)
+@click.option("-a", "--auto-threshold", type=click.FLOAT, default=None)
+@click.option("-l", "--limit", type=click.INT, default=5000)
+@click.option("--algorithm", default=DefaultAlgorithm.NAME)
+@click.option("--scored/--unscored", is_flag=True, type=click.BOOL, default=True)
+@click.option(
+    "-c",
+    "--clear",
+    is_flag=True,
+    default=False,
+    help="Clear the index directory, if it exists.",
+)
+def xref_file(
+    path: Path,
+    auto_threshold: Optional[float] = None,
+    algorithm: str = DefaultAlgorithm.NAME,
+    limit: int = 5000,
+    scored: bool = True,
+    clear: bool = False,
+) -> None:
+    resolver = Resolver[Entity].make_default()
+    resolver.begin()
+    store = load_entity_file_store(path, resolver=resolver)
+    algorithm_type = get_algorithm(algorithm)
+    if algorithm_type is None:
+        raise click.Abort(f"Unknown algorithm: {algorithm}")
+
+    index_dir = Path(
+        os.environ.get("NOMENKLATURA_INDEX_PATH", path.parent / INDEX_SEGMENT)
+    )
+    if clear and index_dir.exists():
+        log.info("Clearing index: %s", index_dir)
+        shutil.rmtree(index_dir, ignore_errors=True)
+    run_xref(
+        resolver,
+        store,
+        index_dir,
+        auto_threshold=auto_threshold,
+        algorithm=algorithm_type,
+        scored=scored,
+        limit=limit,
+    )
+    resolver.commit()
+    log.info("Xref complete in: %r", resolver)
+
+
+@cli.command("prune", help="Remove dedupe candidates")
+def xref_prune() -> None:
+    resolver = Resolver[Entity].make_default()
+    resolver.begin()
+    resolver.prune()
+    resolver.commit()
+
+
+@cli.command("apply", help="Apply resolver to an entity stream")
+@click.argument("path", type=InPath)
+@click.option("-o", "--outpath", type=OutPath, default="-")
+@click.option(
+    "-d",
+    "--dataset",
+    type=str,
+    default=None,
+    help="Add a dataset to the entity metadata",
+)
+def apply(path: Path, outpath: Path, dataset: Optional[str] = None) -> None:
+    linker = _get_linker()
+    with path_writer(outpath) as outfh:
+        for proxy in path_entities(path, ValueEntity):
+            proxy = linker.apply_stream(proxy)
+            if dataset is not None:
+                proxy.datasets.add(dataset)
+            write_entity(outfh, proxy)
+
+
+@cli.command("sorted-aggregate", help="Merge sort-order entities")
+@click.option("-i", "--infile", type=InPath, default="-")
+@click.option("-o", "--outfile", type=OutPath, default="-")
+def sorted_aggregate_(infile: Path, outfile: Path) -> None:
+    sorted_aggregate(infile, outfile, ValueEntity)
+
+
+@cli.command("make-sortable", help="Convert entities into plain-text sortable form")
+@click.argument("path", type=InPath)
+@click.option("-o", "--outpath", type=OutPath, default="-")
+def make_sortable(path: Path, outpath: Path) -> None:
+    with path_writer(outpath) as outfh:
+        for entity in path_entities(path, Entity):
+            write_entity(outfh, entity)
+
+
+@cli.command("dedupe", help="Interactively judge xref candidates")
+@click.argument("path", type=InPath)
+@click.option("-x", "--xref", is_flag=True, default=False)
+def dedupe(path: Path, xref: bool = False) -> None:
+    resolver = Resolver[Entity].make_default()
+    resolver.begin()
+    store = load_entity_file_store(path, resolver=resolver)
+    if xref:
+        index_dir = path.parent / INDEX_SEGMENT
+        run_xref(resolver, store, index_dir)
+    resolver.commit()
+
+    dedupe_ui(resolver, store)
+
+
+@cli.command("train-v1-matcher", help="Train a matching model from judgement pairs")
+@click.argument("pairs_file", type=InPath)
+def train_v1_matcher_(pairs_file: Path) -> None:
+    train_v1_matcher(pairs_file)
+
+
+@cli.command("train-erun-matcher", help="Train an ER model from judgement pairs")
+@click.argument("pairs_file", type=InPath)
+def train_erun_matcher_(pairs_file: Path) -> None:
+    train_erun_matcher(pairs_file)
+
+
+@cli.command("match", help="Generate matches from an enrichment source")
+@click.argument("config", type=InPath)
+@click.argument("entities", type=InPath)
+@click.option("-o", "--outpath", type=OutPath, default="-")
+def match_command(
+    config: Path,
+    entities: Path,
+    outpath: Path,
+) -> None:
+    resolver = Resolver[Entity].make_default()
+    _, enricher = _load_enricher(config)
+
+    try:
+        resolver.begin()
+        with path_writer(outpath) as fh:
+            stream = path_entities(entities, Entity)
+            for proxy in match(enricher, resolver, stream):
+                write_entity(fh, proxy)
+        resolver.commit()
+    finally:
+        enricher.close()
+
+
+@cli.command("enrich", help="Fetch extra info from an enrichment source")
+@click.argument("config", type=InPath)
+@click.argument("entities", type=InPath)
+@click.option("-o", "--outpath", type=OutPath, default="-")  # noqa
+def enrich_command(
+    config: Path,
+    entities: Path,
+    outpath: Path,
+) -> None:
+    resolver = Resolver[Entity].make_default()
+    _, enricher = _load_enricher(config)
+    try:
+        resolver.begin()
+        with path_writer(outpath) as fh:
+            stream = path_entities(entities, Entity)
+            for proxy in enrich(enricher, resolver, stream):
+                write_entity(fh, proxy)
+        resolver.commit()
+    finally:
+        enricher.close()
+
+
+@cli.command("apply-statements", help="Apply a resolver file to a set of statements")
+@click.option("-i", "--infile", type=InPath, default="-")
+@click.option("-o", "--outpath", type=OutPath, default="-")
+@click.option("-f", "--format", type=click.Choice(FORMATS), default=CSV)
+def statements_apply(infile: Path, outpath: Path, format: str) -> None:
+    linker = _get_linker()
+
+    def _generate() -> Generator[Statement, None, None]:
+        for stmt in read_path_statements(infile, format=format):
+            yield linker.apply_statement(stmt)
+
+    with path_writer(outpath) as outfh:
+        write_statements(outfh, format, _generate())
+
+
+@cli.command("load-resolver", help="Load resolver edges from file into database")
+@click.argument("source", type=InPath)
+def load_resolver(source: Path) -> None:
+    resolver = Resolver[Entity].make_default()
+    resolver.begin()
+    resolver.load(source)
+    resolver.commit()
+
+
+@cli.command("dump-resolver", help="Dump resolver decisions from database to file")
+@click.argument("target", type=OutPath)
+def dump_resolver(target: Path) -> None:
+    resolver = Resolver[Entity].make_default()
+    resolver.begin()
+    resolver.dump(target)
+    resolver.rollback()
+
+
+@cli.command("bench", help="Benchmark a matching algorithm")
+@click.argument("name", type=str)
+@click.argument("pairs_file", type=InPath)
+@click.option("-n", "--number", type=int, default=1000)
+def bench(name: str, pairs_file: Path, number: int = 1000) -> None:
+    bench_matcher(name, pairs_file, number)
+
+
+if __name__ == "__main__":
+    cli()

nomenklatura/conflicting_match.py

@@ -0,0 +1,80 @@
+from typing import Dict, Set, Tuple, Generic, Generator
+from itertools import combinations
+from collections import defaultdict
+from rich.console import Console
+from rich.table import Table
+from rich import box
+
+from followthemoney import DS, Statement, SE
+from nomenklatura.store import View
+from nomenklatura.judgement import Judgement
+from nomenklatura.resolver import Resolver
+
+
+class ConflictingMatchReporter(Generic[SE]):
+    def __init__(self, view: View[DS, SE], resolver: Resolver[SE], threshold: float):
+        self.console = Console()
+        self.view = view
+        self.resolver = resolver
+        self.threshold = threshold
+        self.matches: Dict[str, Set[str]] = defaultdict(set)
+
+    def check_match(self, score: float, left_id: str, right_id: str) -> None:
+        if score > self.threshold:
+            self.matches[left_id].add(right_id)
+            self.matches[right_id].add(left_id)
+
+    def get_conflicting_matches(self) -> Generator[Tuple[str, str, str], None, None]:
+        for candidate_id, matches in self.matches.items():
+            for left_id, right_id in combinations(matches, 2):
+                judgement = self.resolver.get_judgement(left_id, right_id)
+                if judgement == Judgement.NEGATIVE:
+                    yield candidate_id, left_id, right_id
+
+    @staticmethod
+    def _sort_key(stmt: Statement) -> Tuple[str, str, int, str, int]:
+        prop_order = 0 if stmt.prop == "name" else 1
+        lang_order = 0 if stmt.lang is None else 1
+        return (stmt.dataset, stmt.entity_id, prop_order, stmt.value, lang_order)
+
+    def report_conflicting_match(self, title: str, entity: SE) -> None:
+        if entity.id is None:
+            return
+        statements = []
+        for stmt in entity.statements:
+            if stmt.prop in {"name", "alias"}:
+                statements.append(stmt)
+
+        table = Table(box=box.SIMPLE, expand=True)
+        table.add_column("Dataset", style="cyan", max_width=20)
+        table.add_column("Entity ID", style="magenta", max_width=30)
+        table.add_column("Prop", style="blue")
+        table.add_column("Lang", style="green")
+        table.add_column("Name", style="yellow")
+
+        for stmt in sorted(statements, key=self._sort_key):
+            table.add_row(
+                stmt.dataset, stmt.entity_id, stmt.prop, stmt.lang, "• " + stmt.value
+            )
+
+        self.console.print(f"[bold]{title}[/bold]:")
+        self.console.print(f"{entity.id}")
+        self.console.print(table)
+
+    def report(self) -> None:
+        conflicts = list(self.get_conflicting_matches())
+        if not conflicts:
+            return
+
+        self.console.print("[bold]Potential conflicting matches found:\n[/bold]")
+        for candidate_id, left_id, right_id in self.get_conflicting_matches():
+            left = self.view.get_entity(left_id)
+            right = self.view.get_entity(right_id)
+            candidate = self.view.get_entity(candidate_id)
+
+            if candidate:
+                self.report_conflicting_match("Candidate", candidate)
+            if left:
+                self.report_conflicting_match("Left side of negative decision", left)
+            if right:
+                self.report_conflicting_match("Right side of negative decision", right)

nomenklatura/db.py

@@ -0,0 +1,139 @@
+from contextlib import contextmanager
+from functools import cache
+from typing import Any, Dict, Generator, Iterable, List, Mapping, Optional, cast
+import logging
+
+from followthemoney import Statement
+from followthemoney.statement.util import get_prop_type
+from sqlalchemy import (
+    Boolean,
+    Column,
+    DateTime,
+    Dialect,
+    MetaData,
+    Table,
+    Unicode,
+    create_engine,
+    delete,
+)
+from sqlalchemy.engine import Connection, Engine
+from sqlalchemy.dialects.postgresql import insert as psql_insert
+from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+from nomenklatura import settings
+
+Conn = Connection
+Connish = Optional[Connection]
+KEY_LEN = 255
+VALUE_LEN = 65535
+
+log = logging.getLogger(__name__)
+
+
+@cache
+def get_engine(url: Optional[str] = None) -> Engine:
+    url = url or settings.DB_URL
+    connect_args = {}
+    if url.startswith("postgres"):
+        connect_args["options"] = f"-c statement_timeout={settings.DB_STMT_TIMEOUT}"
+
+    return create_engine(
+        url, pool_size=settings.DB_POOL_SIZE, connect_args=connect_args
+    )
+
+
+@cache
+def get_metadata() -> MetaData:
+    return MetaData()
+
+
+@contextmanager
+def ensure_tx(conn: Connish = None) -> Generator[Connection, None, None]:
+    if conn is not None:
+        yield conn
+        return
+    engine = get_engine()
+    with engine.begin() as conn:
+        yield conn
+
+
+def make_statement_table(
+    metadata: MetaData,
+    name: str = settings.STATEMENT_TABLE,
+) -> Table:
+    return Table(
+        name,
+        metadata,
+        Column("id", Unicode(KEY_LEN), primary_key=True, unique=True),
+        Column("entity_id", Unicode(KEY_LEN), index=True, nullable=False),
+        Column("canonical_id", Unicode(KEY_LEN), index=True, nullable=False),
+        Column("prop", Unicode(KEY_LEN), index=True, nullable=False),
+        Column("prop_type", Unicode(KEY_LEN), index=True, nullable=False),
+        Column("schema", Unicode(KEY_LEN), index=True, nullable=False),
+        Column("value", Unicode(VALUE_LEN), nullable=False),
+        Column("original_value", Unicode(VALUE_LEN), nullable=True),
+        Column("dataset", Unicode(KEY_LEN), index=True),
+        Column("origin", Unicode(KEY_LEN), index=True),
+        Column("lang", Unicode(KEY_LEN), nullable=True),
+        Column("external", Boolean, default=False, nullable=False),
+        Column("first_seen", DateTime, nullable=True),
+        Column("last_seen", DateTime, nullable=True),
+    )
+
+
+def _upsert_statement_batch(
+    dialect: Dialect, conn: Connection, table: Table, batch: List[Mapping[str, Any]]
+) -> None:
+    """Create an upsert statement for the given table and engine."""
+    if dialect.name == "sqlite":
+        lstmt = sqlite_insert(table).values(batch)
+        lstmt = lstmt.on_conflict_do_nothing(index_elements=["id"])
+        conn.execute(lstmt)
+    elif dialect.name in ("postgresql", "postgres"):
+        pstmt = psql_insert(table).values(batch)
+        pstmt = pstmt.on_conflict_do_nothing(index_elements=["id"])
+        conn.execute(pstmt)
+    else:
+        raise NotImplementedError(f"Upsert not implemented for dialect {dialect.name}")
+
+
+def insert_statements(
+    engine: Engine,
+    table: Table,
+    dataset_name: str,
+    statements: Iterable[Statement],
+    batch_size: int = settings.STATEMENT_BATCH,
+) -> None:
+    dataset_count: int = 0
+    is_postgresql = "postgres" in engine.dialect.name
+    with engine.begin() as conn:
+        del_q = delete(table).where(table.c.dataset == dataset_name)
+        conn.execute(del_q)
+        batch: List[Mapping[str, Any]] = []
+
+        for stmt in statements:
+            if is_postgresql:
+                row = cast(Dict[str, Any], stmt.to_dict())
+                row["prop_type"] = get_prop_type(row["schema"], row["prop"])
+            else:
+                row = stmt.to_db_row()
+            batch.append(row)
+            dataset_count += 1
+            if len(batch) >= batch_size:
+                args = (len(batch), dataset_count, dataset_name)
+                log.info("Inserting batch %s statements (total: %s) into %r" % args)
+                _upsert_statement_batch(engine.dialect, conn, table, batch)
+                batch = []
+        if len(batch):
+            _upsert_statement_batch(engine.dialect, conn, table, batch)
+        log.info("Load complete: %r (%d total)" % (dataset_name, dataset_count))
+
+
+# TODO: consider offering a COPY-based loader:
+# raw_conn = await conn.get_raw_connection()
+# driver_conn: Connection = raw_conn.driver_connection
+# result = await driver_conn.copy_records_to_table(
+#     stmt_table.name,
+#     records=load_data_rows(),
+#     columns=COLUMNS,
+# )

nomenklatura/delta.py

@@ -0,0 +1,4 @@
+ADD = "ADD"
+MOD = "MOD"
+DEL = "DEL"
+OPS = [ADD, MOD, DEL]