followthemoney 3.8.4__py3-none-any.whl → 4.0.0__py3-none-any.whl

followthemoney/statement/serialize.py

@@ -0,0 +1,245 @@
+import csv
+import click
+import orjson
+from io import TextIOWrapper
+from pathlib import Path
+from types import TracebackType
+from typing import cast
+from typing import BinaryIO, Generator, Iterable, List, Optional, TextIO, Type
+from rigour.boolean import text_bool
+
+from followthemoney.statement.statement import Statement, StatementDict
+from followthemoney.statement.util import unpack_prop
+
+
+JSON = "json"
+CSV = "csv"
+PACK = "pack"
+FORMATS = [JSON, CSV, PACK]
+
+CSV_BATCH = 5000
+CSV_COLUMNS = [
+    "canonical_id",
+    "entity_id",
+    "prop",
+    "prop_type",
+    "schema",
+    "value",
+    "dataset",
+    "origin",
+    "lang",
+    "original_value",
+    "external",
+    "first_seen",
+    "last_seen",
+    "id",
+]
+LEGACY_PACK_COLUMNS = [
+    "entity_id",
+    "prop",
+    "value",
+    "dataset",
+    "lang",
+    "original_value",
+    "target",
+    "external",
+    "first_seen",
+    "last_seen",
+]
+
+
+def read_json_statements(
+    fh: BinaryIO,
+    max_line: int = 40 * 1024 * 1024,
+) -> Generator[Statement, None, None]:
+    while line := fh.readline(max_line):
+        data = orjson.loads(line)
+        yield Statement.from_dict(data)
+
+
+def read_csv_statements(fh: BinaryIO) -> Generator[Statement, None, None]:
+    wrapped = TextIOWrapper(fh, encoding="utf-8")
+    for row in csv.DictReader(wrapped, dialect=csv.unix_dialect):
+        data = cast(StatementDict, row)
+        data["external"] = text_bool(row.get("external")) or False
+        if row.get("lang") == "":
+            data["lang"] = None
+        if row.get("original_value") == "":
+            data["original_value"] = None
+        yield Statement.from_dict(data)
+
+
+def read_pack_statements(fh: BinaryIO) -> Generator[Statement, None, None]:
+    wrapped = TextIOWrapper(fh, encoding="utf-8")
+    yield from read_pack_statements_decoded(wrapped)
+
+
+def read_pack_statements_decoded(fh: TextIO) -> Generator[Statement, None, None]:
+    headers: Optional[List[str]] = None
+    for row in csv.reader(fh, dialect=csv.unix_dialect):
+        if headers is None:
+            if "entity_id" in row and "prop" in row:
+                headers = row
+            else:
+                # This is a legacy pack file, with no headers.
+                headers = LEGACY_PACK_COLUMNS
+            continue
+        data = dict(zip(headers, row))
+        schema, _, prop = unpack_prop(data["prop"])
+        yield Statement(
+            entity_id=data["entity_id"],
+            prop=prop,
+            schema=schema,
+            value=data["value"],
+            dataset=data["dataset"],
+            lang=data.get("lang") or None,
+            original_value=data.get("original_value") or None,
+            origin=data.get("origin"),
+            first_seen=data["first_seen"],
+            external=data["external"] == "t",
+            canonical_id=data["entity_id"],
+            last_seen=data["last_seen"],
+            id=data.get("id"),
+        )
+
+
+def read_statements(fh: BinaryIO, format: str) -> Generator[Statement, None, None]:
+    if format == CSV:
+        yield from read_csv_statements(fh)
+    elif format == PACK:
+        yield from read_pack_statements(fh)
+    else:
+        yield from read_json_statements(fh)
+
+
+def read_path_statements(path: Path, format: str) -> Generator[Statement, None, None]:
+    if str(path) == "-":
+        fh = click.get_binary_stream("stdin")
+        yield from read_statements(fh, format=format)
+        return
+    with open(path, "rb") as fh:
+        yield from read_statements(fh, format=format)
+
+
+def get_statement_writer(fh: BinaryIO, format: str) -> "StatementWriter":
+    if format == CSV:
+        wrapped = TextIOWrapper(fh, encoding="utf-8")
+        return CSVStatementWriter(wrapped)
+    elif format == PACK:
+        wrapped = TextIOWrapper(fh, encoding="utf-8")
+        return PackStatementWriter(wrapped)
+    elif format == JSON:
+        return JSONStatementWriter(fh)
+    raise RuntimeError("Unknown statement format: %s" % format)
+
+
+def write_statements(
+    fh: BinaryIO, format: str, statements: Iterable[Statement]
+) -> None:
+    writer = get_statement_writer(fh, format)
+    for stmt in statements:
+        writer.write(stmt)
+    writer.close()
+
+
+class StatementWriter(object):
+    def write(self, stmt: Statement) -> None:
+        raise NotImplementedError()
+
+    def close(self) -> None:
+        raise NotImplementedError()
+
+    def __enter__(self) -> "StatementWriter":
+        return self
+
+    def __exit__(
+        self,
+        type: Optional[Type[BaseException]],
+        value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
+        self.close()
+
+
+class JSONStatementWriter(StatementWriter):
+    def __init__(self, fh: BinaryIO) -> None:
+        self.fh = fh
+
+    def write(self, stmt: Statement) -> None:
+        data = stmt.to_dict()
+        out = orjson.dumps(data, option=orjson.OPT_APPEND_NEWLINE)
+        self.fh.write(out)
+
+    def close(self) -> None:
+        self.fh.close()
+
+
+class CSVStatementWriter(StatementWriter):
+    def __init__(self, fh: TextIO) -> None:
+        self.fh = fh
+        self.writer = csv.writer(self.fh, dialect=csv.unix_dialect)
+        self.writer.writerow(CSV_COLUMNS)
+        self._batch: List[List[Optional[str]]] = []
+
+    def write(self, stmt: Statement) -> None:
+        row = stmt.to_csv_row()
+        self._batch.append([row[c] for c in CSV_COLUMNS])
+        if len(self._batch) >= CSV_BATCH:
+            self.writer.writerows(self._batch)
+            self._batch.clear()
+
+    def close(self) -> None:
+        if len(self._batch) > 0:
+            self.writer.writerows(self._batch)
+        self.fh.close()
+
+
+class PackStatementWriter(StatementWriter):
+    def __init__(self, fh: TextIO) -> None:
+        self.fh = fh
+        self.writer = csv.writer(
+            self.fh,
+            dialect=csv.unix_dialect,
+            quoting=csv.QUOTE_MINIMAL,
+        )
+        columns = [
+            "entity_id",
+            "prop",
+            "value",
+            "dataset",
+            "lang",
+            "original_value",
+            "origin",
+            "external",
+            "first_seen",
+            "last_seen",
+            "id",
+        ]
+        self.writer.writerow(columns)
+        self._batch: List[List[Optional[str]]] = []
+
+    def write(self, stmt: Statement) -> None:
+        # HACK: This is very similar to the CSV writer, but at the very inner
+        # loop of the application, so we're duplicating code here.
+        row = [
+            stmt.entity_id,
+            f"{stmt.schema}:{stmt.prop}",
+            stmt.value,
+            stmt.dataset,
+            stmt.lang,
+            stmt.original_value,
+            stmt.origin,
+            "t" if stmt.external else None,
+            stmt.first_seen,
+            stmt.last_seen,
+            stmt.id,
+        ]
+        self._batch.append(row)
+        if len(self._batch) >= CSV_BATCH:
+            self.writer.writerows(self._batch)
+            self._batch.clear()
+
+    def close(self) -> None:
+        if len(self._batch) > 0:
+            self.writer.writerows(self._batch)
+        self.fh.close()

followthemoney/statement/statement.py

@@ -0,0 +1,256 @@
+import hashlib
+import warnings
+from sqlalchemy.engine import Row
+from typing import cast
+from typing import Any, Dict, Generator, Optional
+from typing_extensions import TypedDict, Self
+from rigour.time import datetime_iso, iso_datetime
+from rigour.boolean import bool_text
+
+from followthemoney.proxy import EntityProxy
+from followthemoney.statement.util import get_prop_type, BASE_ID
+
+
+class StatementDict(TypedDict):
+    id: Optional[str]
+    entity_id: str
+    canonical_id: str
+    prop: str
+    schema: str
+    value: str
+    dataset: str
+    lang: Optional[str]
+    original_value: Optional[str]
+    external: bool
+    first_seen: Optional[str]
+    last_seen: Optional[str]
+    origin: Optional[str]
+
+
+class Statement(object):
+    """A single statement about a property relevant to an entity.
+
+    For example, this could be used to say: "In dataset A, entity X has the
+    property `name` set to 'John Smith'. I first observed this at K, and last
+    saw it at L."
+
+    Null property values are not supported. This might need to change if we
+    want to support making property-less entities.
+    """
+
+    BASE = BASE_ID
+
+    __slots__ = [
+        "id",
+        "entity_id",
+        "canonical_id",
+        "prop",
+        "schema",
+        "value",
+        "dataset",
+        "lang",
+        "original_value",
+        "external",
+        "first_seen",
+        "last_seen",
+        "origin",
+    ]
+
+    def __init__(
+        self,
+        entity_id: str,
+        prop: str,
+        schema: str,
+        value: str,
+        dataset: str,
+        lang: Optional[str] = None,
+        original_value: Optional[str] = None,
+        first_seen: Optional[str] = None,
+        external: bool = False,
+        id: Optional[str] = None,
+        canonical_id: Optional[str] = None,
+        last_seen: Optional[str] = None,
+        origin: Optional[str] = None,
+    ):
+        self.entity_id = entity_id
+        self.canonical_id = canonical_id or entity_id
+        self.prop = prop
+        self.schema = schema
+        self.value = value
+        self.dataset = dataset
+        self.lang = lang
+        self.original_value = original_value
+        self.first_seen = first_seen
+        self.last_seen = last_seen or first_seen
+        self.external = external
+        self.origin = origin
+        if id is None:
+            id = self.generate_key()
+        self.id = id
+
+    @property
+    def prop_type(self) -> str:
+        """The type of the property, e.g. 'string', 'number', 'url'."""
+        return get_prop_type(self.schema, self.prop)
+
+    def to_dict(self) -> StatementDict:
+        return {
+            "canonical_id": self.canonical_id,
+            "entity_id": self.entity_id,
+            "prop": self.prop,
+            "schema": self.schema,
+            "value": self.value,
+            "dataset": self.dataset,
+            "lang": self.lang,
+            "original_value": self.original_value,
+            "first_seen": self.first_seen,
+            "last_seen": self.last_seen,
+            "external": self.external,
+            "origin": self.origin,
+            "id": self.id,
+        }
+
+    def to_csv_row(self) -> Dict[str, Optional[str]]:
+        data = cast(Dict[str, Optional[str]], self.to_dict())
+        data["external"] = bool_text(self.external)
+        data["prop_type"] = get_prop_type(self.schema, self.prop)
+        return data
+
+    def to_db_row(self) -> Dict[str, Any]:
+        data = cast(Dict[str, Any], self.to_dict())
+        data["first_seen"] = iso_datetime(self.first_seen)
+        data["last_seen"] = iso_datetime(self.last_seen)
+        data["prop_type"] = get_prop_type(self.schema, self.prop)
+        return data
+
+    def __hash__(self) -> int:
+        if self.id is None:
+            warnings.warn(
+                "Hashing a statement without an ID results in undefined behaviour",
+                RuntimeWarning,
+            )
+        return hash(self.id)
+
+    def __repr__(self) -> str:
+        return "<Statement(%r, %r, %r)>" % (self.entity_id, self.prop, self.value)
+
+    def __eq__(self, other: Any) -> bool:
+        return not self.id != other.id
+
+    def __lt__(self, other: Any) -> bool:
+        self_key = (self.prop != BASE_ID, self.id or "")
+        other_key = (other.prop != BASE_ID, other.id or "")
+        return self_key < other_key
+
+    def clone(self: Self) -> "Statement":
+        """Make a deep copy of the given statement."""
+        return Statement.from_dict(self.to_dict())
+
+    def generate_key(self) -> Optional[str]:
+        return self.make_key(
+            self.dataset,
+            self.entity_id,
+            self.prop,
+            self.value,
+            self.external,
+        )
+
+    @classmethod
+    def make_key(
+        cls,
+        dataset: str,
+        entity_id: str,
+        prop: str,
+        value: str,
+        external: Optional[bool],
+    ) -> Optional[str]:
+        """Hash the key properties of a statement record to make a unique ID."""
+        if prop is None or value is None:
+            return None
+        key = f"{dataset}.{entity_id}.{prop}.{value}"
+        if external:
+            # We consider the external flag in key composition to avoid race conditions
+            # where a certain entity might be emitted as external while it is already
+            # linked in to the graph via another route.
+            key = f"{key}.ext"
+        return hashlib.sha1(key.encode("utf-8")).hexdigest()
+
+    @classmethod
+    def from_dict(cls, data: StatementDict) -> "Statement":
+        return cls(
+            entity_id=data["entity_id"],
+            prop=data["prop"],
+            schema=data["schema"],
+            value=data["value"],
+            dataset=data["dataset"],
+            lang=data.get("lang", None),
+            original_value=data.get("original_value", None),
+            first_seen=data.get("first_seen", None),
+            external=data.get("external", False),
+            id=data.get("id", None),
+            canonical_id=data.get("canonical_id", None),
+            last_seen=data.get("last_seen", None),
+            origin=data.get("origin", None),
+        )
+
+    @classmethod
+    def from_db_row(cls, row: Row[Any]) -> "Statement":
+        return cls(
+            id=row.id,
+            canonical_id=row.canonical_id,
+            entity_id=row.entity_id,
+            prop=row.prop,
+            schema=row.schema,
+            value=row.value,
+            dataset=row.dataset,
+            lang=row.lang,
+            original_value=row.original_value,
+            first_seen=datetime_iso(row.first_seen),
+            external=row.external,
+            last_seen=datetime_iso(row.last_seen),
+            origin=row.origin,
+        )
+
+    @classmethod
+    def from_entity(
+        cls,
+        entity: "EntityProxy",
+        dataset: str,
+        first_seen: Optional[str] = None,
+        last_seen: Optional[str] = None,
+        external: bool = False,
+        origin: Optional[str] = None,
+    ) -> Generator["Statement", None, None]:
+        from followthemoney.statement.entity import StatementEntity
+
+        if entity.id is None:
+            raise ValueError("Cannot create statements for entity without ID!")
+
+        # If the entity is already a StatementEntity, we return its statements directly.
+        if isinstance(entity, StatementEntity):
+            yield from entity.statements
+            return
+
+        yield cls(
+            entity_id=entity.id,
+            prop=BASE_ID,
+            schema=entity.schema.name,
+            value=entity.id,
+            dataset=dataset,
+            external=external,
+            first_seen=first_seen,
+            last_seen=last_seen,
+            origin=origin,
+        )
+        for prop, value in entity.itervalues():
+            yield cls(
+                entity_id=entity.id,
+                prop=prop.name,
+                schema=entity.schema.name,
+                value=value,
+                dataset=dataset,
+                external=external,
+                first_seen=first_seen,
+                last_seen=last_seen,
+                origin=origin,
+            )

followthemoney/statement/util.py

@@ -0,0 +1,31 @@
+import sys
+from functools import cache
+from typing import Tuple
+
+from followthemoney.model import Model
+
+BASE_ID = "id"
+
+
+def pack_prop(schema: str, prop: str) -> str:
+    return f"{schema}:{prop}"
+
+
+@cache
+def get_prop_type(schema: str, prop: str) -> str:
+    if prop == BASE_ID:
+        return BASE_ID
+    schema_obj = Model.instance().get(schema)
+    if schema_obj is None:
+        raise TypeError("Schema not found: %s" % schema)
+    prop_obj = schema_obj.get(prop)
+    if prop_obj is None:
+        raise TypeError("Property not found: %s" % prop)
+    return prop_obj.type.name
+
+
+@cache
+def unpack_prop(id: str) -> Tuple[str, str, str]:
+    schema, prop = id.split(":", 1)
+    prop_type = get_prop_type(schema, prop)
+    return sys.intern(schema), prop_type, sys.intern(prop)

followthemoney/types/__init__.py

@@ -1,4 +1,6 @@
-from followthemoney.types.registry import Registry
+from banal import ensure_list
+from typing import Dict, Iterable, List, Set, cast
+
 from followthemoney.types.url import UrlType
 from followthemoney.types.name import NameType
 from followthemoney.types.email import EmailType
@@ -11,7 +13,6 @@ from followthemoney.types.language import LanguageType
 from followthemoney.types.mimetype import MimeType
 from followthemoney.types.checksum import ChecksumType
 from followthemoney.types.identifier import IdentifierType
-from followthemoney.types.iban import IbanType
 from followthemoney.types.entity import EntityType
 from followthemoney.types.topic import TopicType
 from followthemoney.types.gender import GenderType
@@ -22,27 +23,69 @@ from followthemoney.types.string import StringType
 from followthemoney.types.number import NumberType
 from followthemoney.types.common import PropertyType
 
+
+class Registry(object):
+    """This registry keeps the processing helpers for all property types in the system. The
+    registry can be used to get a type, which can itself then clean, validate or format values
+    of that type."""
+
+    url = UrlType()
+    name = NameType()
+    email = EmailType()
+    ip = IpType()
+    address = AddressType()
+    date = DateType()
+    phone = PhoneType()
+    country = CountryType()
+    language = LanguageType()
+    mimetype = MimeType()
+    checksum = ChecksumType()
+    identifier = IdentifierType()
+    entity = EntityType()
+    topic = TopicType()
+    gender = GenderType()
+    json = JsonType()
+    text = TextType()
+    html = HTMLType()
+    string = StringType()
+    number = NumberType()
+
+    def __init__(self) -> None:
+        self.matchable: Set[PropertyType] = set()
+        self.types: Set[PropertyType] = set()
+        self.groups: Dict[str, PropertyType] = {}
+        self.pivots: Set[PropertyType] = set()
+        for name in dir(self):
+            type_ = getattr(self, name)
+            if not isinstance(type_, PropertyType):
+                continue
+            assert type_.name == name
+            self.types.add(type_)
+            if type_.matchable:
+                self.matchable.add(type_)
+            if type_.pivot:
+                self.pivots.add(type_)
+            if type_.group is not None:
+                self.groups[type_.group] = type_
+
+    def get(self, name: str) -> PropertyType:
+        """For a given property type name, get its type object. This can also
+        be used via getattr, e.g. ``registry.phone``."""
+        # Allow transparent re-checking.
+        if isinstance(name, PropertyType):
+            return name
+        return cast(PropertyType, getattr(self, name))
+
+    def get_types(self, names: Iterable[str]) -> List[PropertyType]:
+        """Get a list of all property type objects linked to a set of names."""
+        names = ensure_list(names)
+        types = [self.get(n) for n in names]
+        return [t for t in types if t is not None]
+
+    def __getitem__(self, name: str) -> PropertyType:
+        return cast(PropertyType, getattr(self, name))
+
+
 registry = Registry()
-registry.add(UrlType)
-registry.add(NameType)
-registry.add(EmailType)
-registry.add(IpType)
-registry.add(AddressType)
-registry.add(DateType)
-registry.add(PhoneType)
-registry.add(CountryType)
-registry.add(LanguageType)
-registry.add(MimeType)
-registry.add(ChecksumType)
-registry.add(IdentifierType)
-registry.add(IbanType)  # TODO: remove
-registry.add(EntityType)
-registry.add(TopicType)
-registry.add(GenderType)
-registry.add(JsonType)
-registry.add(TextType)
-registry.add(HTMLType)
-registry.add(StringType)
-registry.add(NumberType)
 
 __all__ = ["PropertyType", "registry"]

followthemoney/types/address.py

@@ -7,7 +7,7 @@ from rigour.text.distance import levenshtein_similarity
 
 from followthemoney.types.common import PropertyType
 from followthemoney.util import defer as _
-from followthemoney.util import dampen
+from followthemoney.util import dampen, const
 
 if TYPE_CHECKING:
     from followthemoney.proxy import EntityProxy
@@ -21,8 +21,8 @@ class AddressType(PropertyType):
 
     LINE_BREAKS = re.compile(r"(\r\n|\n|<BR/>|<BR>|\t|ESQ\.,|ESQ,|;)")
     COMMATA = re.compile(r"(,\s?[,\.])")
-    name = "address"
-    group = "addresses"
+    name = const("address")
+    group = const("addresses")
     label = _("Address")
     plural = _("Addresses")
     matchable = True

followthemoney/types/checksum.py

@@ -1,6 +1,5 @@
-from followthemoney.rdf import URIRef, Identifier
 from followthemoney.types.common import PropertyType
-from followthemoney.util import defer as _
+from followthemoney.util import const, defer as _
 
 
 class ChecksumType(PropertyType):
@@ -13,13 +12,10 @@ class ChecksumType(PropertyType):
     of this type are scrubbed when submitted via the normal API. Checksums can only
     be defined by uploading a document to be ingested."""
 
-    name = "checksum"
-    group = "checksums"
+    name = const("checksum")
+    group = const("checksums")
     label = _("Checksum")
     plural = _("Checksums")
     matchable = True
     pivot = True
     max_length = 40
-
-    def rdf(self, value: str) -> Identifier:
-        return URIRef(f"hash:{value}")