ftmq 4.3.1__py3-none-any.whl → 4.3.4__py3-none-any.whl

ftmq/__init__.py

@@ -2,7 +2,7 @@ from ftmq.io import smart_read_proxies, smart_write_proxies
 from ftmq.query import Query
 from ftmq.util import make_entity
 
-__version__ = "4.3.1"
+__version__ = "4.3.4"
 __all__ = [
     "smart_read_proxies",
     "smart_write_proxies",

ftmq/model/dataset.py

@@ -16,6 +16,7 @@ ContentType = Literal["documents", "structured", "mixed"]
 
 
 class Dataset(BaseModel, _DatasetModel):
+    prefix: str | None = None
     maintainer: DataPublisher | None = None
     stats: DatasetStats = DatasetStats()
     git_repo: AnyUrl | None = None

ftmq/store/lake.py

@@ -19,7 +19,7 @@ Layout:
 """
 
 from pathlib import Path
-from typing import Any, Generator, Iterable
+from typing import Any, Generator
 from urllib.parse import urlparse
 
 import duckdb
@@ -62,6 +62,9 @@ Z_ORDER = ["canonical_id", "entity_id", "schema", "prop"]
 TARGET_SIZE = 50 * 10_485_760  # 500 MB
 PARTITION_BY = ["dataset", "bucket", "origin"]
 DEFAULT_ORIGIN = "default"
+BUCKET_MENTION = "mention"
+BUCKET_PAGE = "page"
+BUCKET_PAGES = "pages"
 BUCKET_DOCUMENT = "document"
 BUCKET_INTERVAL = "interval"
 BUCKET_THING = "thing"
@@ -93,6 +96,7 @@ TABLE = table(
     column("dataset"),
     column("bucket"),
     column("origin"),
+    column("source"),
     column("schema"),
     column("prop"),
     column("prop_type"),
@@ -164,6 +168,12 @@ def setup_duckdb_storage() -> None:
 @cache
 def get_schema_bucket(schema_name: str) -> str:
     s = model[schema_name]
+    if s.is_a("Mention"):
+        return BUCKET_MENTION
+    if s.is_a("Page"):
+        return BUCKET_PAGE
+    if s.is_a("Pages"):
+        return BUCKET_PAGES
     if s.is_a("Document"):
         return BUCKET_DOCUMENT
     if s.is_a("Interval"):
@@ -171,19 +181,13 @@ def get_schema_bucket(schema_name: str) -> str:
     return BUCKET_THING
 
 
-def pack_statement(stmt: Statement) -> SDict:
+def pack_statement(stmt: Statement, source: str | None = None) -> SDict:
     data = stmt.to_db_row()
     data["bucket"] = get_schema_bucket(data["schema"])
+    data["source"] = source
     return data
 
 
-def pack_statements(statements: Iterable[Statement]) -> pd.DataFrame:
-    df = pd.DataFrame(map(pack_statement, statements))
-    df = df.drop_duplicates()  # .sort_values(Z_ORDER)
-    df = df.fillna(np.nan)
-    return df
-
-
 def compile_query(q: Select) -> str:
     table = nks.STATEMENT_TABLE
     sql = str(q.compile(compile_kwargs={"literal_binds": True}))
@@ -243,9 +247,6 @@ class LakeStore(SQLStore):
         self._partition_by = kwargs.pop("partition_by", PARTITION_BY)
         self._lock: Lock = kwargs.pop("lock", Lock(self._backend))
         self._enforce_dataset = kwargs.pop("enforce_dataset", False)
-        assert isinstance(
-            self._backend, FSStore
-        ), f"Invalid store backend: `{self._backend.__class__}"
         kwargs["uri"] = "sqlite:///:memory:"  # fake it till you make it
         get_metadata.cache_clear()
         super().__init__(*args, **kwargs)
@@ -279,8 +280,10 @@ class LakeStore(SQLStore):
         scope = scope or self.dataset
         return LakeQueryView(self, scope, external)
 
-    def writer(self, origin: str | None = DEFAULT_ORIGIN) -> "LakeWriter":
-        return LakeWriter(self, origin=origin or DEFAULT_ORIGIN)
+    def writer(
+        self, origin: str | None = DEFAULT_ORIGIN, source: str | None = None
+    ) -> "LakeWriter":
+        return LakeWriter(self, origin=origin or DEFAULT_ORIGIN, source=source)
 
     def get_origins(self) -> set[str]:
         q = select(self.table.c.origin).distinct()
@@ -291,32 +294,50 @@ class LakeWriter(nk.Writer):
     store: LakeStore
     BATCH_STATEMENTS = 1_000_000
 
-    def __init__(self, store: Store, origin: str | None = DEFAULT_ORIGIN):
+    def __init__(
+        self,
+        store: Store,
+        origin: str | None = DEFAULT_ORIGIN,
+        source: str | None = None,
+    ):
         super().__init__(store)
-        self.batch: set[Statement] = set()
+        self.batch: dict[Statement, str | None] = {}
         self.origin = origin or DEFAULT_ORIGIN
+        self.source = source
 
-    def add_statement(self, stmt: Statement) -> None:
+    def add_statement(self, stmt: Statement, source: str | None = None) -> None:
         if stmt.entity_id is None:
             return
         stmt.origin = stmt.origin or self.origin
         canonical_id = self.store.linker.get_canonical(stmt.entity_id)
         stmt.canonical_id = canonical_id
-        self.batch.add(stmt)
+        self.batch[stmt] = source or self.source
 
-    def add_entity(self, entity: EntityProxy, origin: str | None = None) -> None:
+    def add_entity(
+        self,
+        entity: EntityProxy,
+        origin: str | None = None,
+        source: str | None = None,
+    ) -> None:
         e = ensure_entity(entity, StatementEntity, self.store.dataset)
         if self.store._enforce_dataset:
             e = apply_dataset(e, self.store.dataset, replace=True)
         for stmt in e.statements:
             if origin:
                 stmt.origin = origin
-            self.add_statement(stmt)
+            self.add_statement(stmt, source=source)
         # we check here instead of in `add_statement` as this will keep entities
-        # together in the same parquet files`
+        # together in the same parquet files
         if len(self.batch) >= self.BATCH_STATEMENTS:
             self.flush()
 
+    def _pack_statements(self) -> pd.DataFrame:
+        data = [pack_statement(stmt, source) for stmt, source in self.batch.items()]
+        df = pd.DataFrame(data)
+        df = df.drop_duplicates().sort_values(Z_ORDER)
+        df = df.fillna(np.nan)
+        return df
+
     def flush(self) -> None:
         if self.batch:
             log.info(
@@ -326,7 +347,7 @@ class LakeWriter(nk.Writer):
             with self.store._lock:
                 write_deltalake(
                     str(self.store.uri),
-                    pack_statements(self.batch),
+                    self._pack_statements(),
                     partition_by=self.store._partition_by,
                     mode="append",
                     schema_mode="merge",
@@ -335,7 +356,7 @@ class LakeWriter(nk.Writer):
                     storage_options=storage_options(),
                 )
 
-        self.batch = set()
+        self.batch = {}
 
     def pop(self, entity_id: str) -> list[Statement]:
         q = select(TABLE)
@@ -353,13 +374,14 @@ class LakeWriter(nk.Writer):
         """
         Optimize the storage: Z-Ordering and compacting
         """
-        self.store.deltatable.optimize.z_order(
-            Z_ORDER, writer_properties=WRITER, target_size=TARGET_SIZE
-        )
-        if vacuum:
-            self.store.deltatable.vacuum(
-                retention_hours=vacuum_keep_hours,
-                enforce_retention_duration=False,
-                dry_run=False,
-                full=True,
+        with self.store._lock:
+            self.store.deltatable.optimize.z_order(
+                Z_ORDER, writer_properties=WRITER, target_size=TARGET_SIZE
             )
+            if vacuum:
+                self.store.deltatable.vacuum(
+                    retention_hours=vacuum_keep_hours,
+                    enforce_retention_duration=False,
+                    dry_run=False,
+                    full=True,
+                )

ftmq/util.py

@@ -1,7 +1,5 @@
-from functools import lru_cache
 from typing import Any, Generator, Type
 
-import pycountry
 from anystore.functools import weakref_cache as cache
 from anystore.types import SDict, StrGenerator
 from banal import ensure_list, is_listish
@@ -13,9 +11,10 @@ from followthemoney.proxy import EntityProxy
 from followthemoney.schema import Schema
 from followthemoney.types import registry
 from followthemoney.util import make_entity_id, sanitize_text
-from normality import collapse_spaces, latinize_text, slugify
+from normality import latinize_text, slugify, squash_spaces
 from rigour.names import Name, Symbol, tag_org_name, tag_person_name
 from rigour.names.tokenize import normalize_name
+from rigour.territories import lookup_territory
 from rigour.text.scripts import can_latinize
 
 from ftmq.enums import Comparators
@@ -149,7 +148,7 @@ def apply_dataset(entity: E, dataset: str | Dataset, replace: bool | None = Fals
 def get_country_name(code: str) -> str:
     """
     Get the (english) country name for the given 2-letter iso code via
-    [pycountry](https://pypi.org/project/pycountry/)
+    [rigour.territories](https://rigour.followthemoney.tech/territories/)
 
     Examples:
         >>> get_country_name("de")
@@ -165,22 +164,17 @@ def get_country_name(code: str) -> str:
     Returns:
         Either the country name for a valid code or the code as fallback.
     """
-    code_clean = get_country_code(code)
-    if code_clean is None:
-        code_clean = code.lower()
-    try:
-        country = pycountry.countries.get(alpha_2=code_clean)
-        if country is not None:
-            return country.name
-    except (LookupError, AttributeError):
-        return code
-    return code_clean
+    territory = lookup_territory(code)
+    if territory is not None:
+        return territory.name
+    return code
 
 
-@lru_cache(1024)
+@cache
 def get_country_code(value: Any, splitter: str | None = ",") -> str | None:
     """
-    Get the 2-letter iso country code for an arbitrary country name
+    Get the 2-letter iso country code for an arbitrary country name via
+    [rigour.territories](https://rigour.followthemoney.tech/territories/)
 
     Examples:
         >>> get_country_code("Germany")
@@ -201,15 +195,16 @@ def get_country_code(value: Any, splitter: str | None = ",") -> str | None:
     """
     value = clean_string(value)
     if not value:
-        return
-    code = registry.country.clean_text(value)
-    if code:
-        return code
-    for token in value.split(splitter):
-        code = registry.country.clean_text(token)
-        if code:
-            return code
-    return
+        return None
+    territory = lookup_territory(value)
+    if territory is not None:
+        return territory.ftm_country
+    if splitter:
+        for token in value.split(splitter):
+            territory = lookup_territory(token.strip())
+            if territory is not None:
+                return territory.ftm_country
+    return None
 
 
 def join_slug(
@@ -321,7 +316,7 @@ def clean_string(value: Any) -> str | None:
     value = sanitize_text(value)
     if value is None:
         return
-    return collapse_spaces(value)
+    return squash_spaces(value)
 
 
 def clean_name(value: Any) -> str | None:

{ftmq-4.3.1.dist-info → ftmq-4.3.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ftmq
-Version: 4.3.1
+Version: 4.3.4
 Summary: followthemoney query dsl and io helpers
 License: AGPLv3+
 License-File: LICENSE
@@ -21,25 +21,24 @@ Provides-Extra: postgres
 Provides-Extra: redis
 Provides-Extra: sql
 Requires-Dist: alephclient (>=2.6.0,<3.0.0) ; extra == "aleph"
-Requires-Dist: anystore (>=0.4.0,<0.5.0)
+Requires-Dist: anystore (>=0.4.2,<0.5.0)
 Requires-Dist: click (>=8.2.1,<9.0.0)
 Requires-Dist: click-default-group (>=1.2.4,<2.0.0)
 Requires-Dist: deltalake (>=1.2.1,<2.0.0) ; extra == "lake"
 Requires-Dist: duckdb (>=1.4.1,<2.0.0) ; extra == "lake"
 Requires-Dist: fakeredis (>=2.26.2,<3.0.0) ; extra == "redis"
-Requires-Dist: followthemoney (>=4.3.2,<5.0.0)
+Requires-Dist: followthemoney (>=4.3.4,<5.0.0)
 Requires-Dist: furl (>=2.1.4,<3.0.0) ; extra == "aleph"
-Requires-Dist: nomenklatura (>=4.1.10,<5.0.0)
+Requires-Dist: nomenklatura (>=4.3.2,<5.0.0)
 Requires-Dist: orjson (>=3.10.18,<4.0.0)
 Requires-Dist: pandas (>=2.3.3,<3.0.0) ; extra == "lake"
 Requires-Dist: plyvel (>=1.5.1,<2.0.0) ; extra == "level"
 Requires-Dist: psycopg[pool] (>=3.2.9,<4.0.0) ; extra == "postgres"
 Requires-Dist: pyarrow (>=21.0.0,<22.0.0) ; extra == "lake"
-Requires-Dist: pycountry (>=24.6.1,<25.0.0)
 Requires-Dist: pydantic (>=2.11.3,<3.0.0)
 Requires-Dist: pyicu (>=2.15.2,<3.0.0)
 Requires-Dist: redis (>=5.2.1,<6.0.0) ; extra == "redis"
-Requires-Dist: rigour (>=1.4.0,<2.0.0)
+Requires-Dist: rigour (>=1.5.0,<2.0.0)
 Requires-Dist: sqlalchemy (>=2.0.36,<3.0.0) ; extra == "postgres"
 Requires-Dist: sqlalchemy (>=2.0.36,<3.0.0) ; extra == "sql"
 Project-URL: Documentation, https://docs.investigraph.dev/lib/ftmq

{ftmq-4.3.1.dist-info → ftmq-4.3.4.dist-info}/RECORD

@@ -1,4 +1,4 @@
-ftmq/__init__.py,sha256=NbZU0yvkkgqiSaOjVHmRAE8FeeI0fIiwknWOfU9j1Ow,245
+ftmq/__init__.py,sha256=7UYwz5cBV9gTOnLRo15i9E2HMsAZFU5DwaqNeo5HAK4,245
 ftmq/aggregate.py,sha256=nyAI5w6jKG1T4Jf2yy1ruhPh0vY6p7JWOEYh0SLBdZY,2163
 ftmq/aggregations.py,sha256=YLu1WF3WgcI3tadWKqsoZk3f_3bYUJetIREy2N1u_EM,4794
 ftmq/cli.py,sha256=e4wLGfGFmn-8Z6_EhEcOJvUwQM3R-J1i40qh2UVHcGo,12078
@@ -7,7 +7,7 @@ ftmq/filters.py,sha256=TKs454wbSvA5tPj8WbpIrMojctE2jGTLnrrVQKM9PE4,7908
 ftmq/io.py,sha256=gUbqoZuYXWwxoiJ505H7FhmXLYr5MEUiCEvqIiOkbgo,3789
 ftmq/logging.py,sha256=p6UvN8qkki8bOjI_NfIjKoLRhdtAdGm4swMWPssETKY,3211
 ftmq/model/__init__.py,sha256=hgdLNeoYm4o6XoY24_7gtIYf8SUq3d8s_2mda1EQDzg,197
-ftmq/model/dataset.py,sha256=Z8RQNz4wvYsbeX77EabI9wH-7difC32_Y_b1QrJkkXA,2198
+ftmq/model/dataset.py,sha256=d6KQDhvxhfQv7J4oF76_ePANDJQ8onxpzhc4-wnui_Q,2228
 ftmq/model/entity.py,sha256=MIJWORs6X4Z22DDQkHRQlcR3Hug8M_YjZwVEo5KU538,2073
 ftmq/model/mixins.py,sha256=O_z3Pxv0HzXxu5K-lNFglrIAhh4jejmbcXUjQfVe0_o,435
 ftmq/model/stats.py,sha256=BiOK7x-JymI5f3dRpwgr3J7OknDP1pe9955Y7AvOJGo,4048
@@ -23,16 +23,16 @@ ftmq/store/fragments/loader.py,sha256=iVh8F22IApe9MRY_Z2fOLvT80fCYstFyxu410l4pPQ
 ftmq/store/fragments/settings.py,sha256=4c-BW-blVM9gC_IGPch03eExbZYFZ3V5h9yTfhcHvOI,303
 ftmq/store/fragments/store.py,sha256=LiSfg95LjEmyq2IUpX4CMtp2tE37SEBLtAw0EWufImM,2534
 ftmq/store/fragments/utils.py,sha256=SDoLPFF5O_oJLIPrCEab5iGn-pl6y0AhYZDYIPYxYkk,1098
-ftmq/store/lake.py,sha256=snbEZXTuR9Oy3o7p6XA61IVpsd4-70a6PgoL63PBrRg,11683
+ftmq/store/lake.py,sha256=PW6uwPeyMo-7bxLUK0zINzOgmo6N1Pulx_Xn-7Y2010,12252
 ftmq/store/level.py,sha256=ZGx-mMtJZJNWkpvbe7ajTREnW5MPcnw0ct3nSFLVF0I,781
 ftmq/store/memory.py,sha256=lZ_pDzrBWNljbNb1MXJeCoO7TnAdqEfG4kfLDOU5rME,551
 ftmq/store/redis.py,sha256=d0hkGF_BezdIfCMUshXWoQwvGmqT8JFblUMcCxzwkDA,433
 ftmq/store/sql.py,sha256=6h3-gDaTAlD-IkiOONcX-JbaAO9-QfSsMjjMPupclcQ,5216
 ftmq/types.py,sha256=HgF8eT3ynKnDUxBYFtoDytS-uN_CS7Yr3DHIX2r4tnk,774
-ftmq/util.py,sha256=CmbZXYAbsKbAjoWn8WxR1Sz4VPXc2gj9CkHwaTqpBG0,15691
-ftmq-4.3.1.dist-info/METADATA,sha256=4sW694l-QjyhMtxl7CN1M2vyFwPDckDg4kpso9aD0cs,5324
-ftmq-4.3.1.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-ftmq-4.3.1.dist-info/entry_points.txt,sha256=YGDCjEiPgAMaQ5MqFKH8m-XIybehSXgarDucSlmeK3E,37
-ftmq-4.3.1.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
-ftmq-4.3.1.dist-info/licenses/NOTICE,sha256=LNgfzuMbk3kIP_KnyDiXO8rQJmDrLy_PQ7cAY8lCmMM,463
-ftmq-4.3.1.dist-info/RECORD,,
+ftmq/util.py,sha256=8N3NVMnfqzOzCpd7tdpeBlWy37sBoLplzud8isX4r8M,15677
+ftmq-4.3.4.dist-info/METADATA,sha256=l7INHEM2JnzmexcV9XOYVqAuYapVx1F1rdb173aRmqE,5279
+ftmq-4.3.4.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+ftmq-4.3.4.dist-info/entry_points.txt,sha256=YGDCjEiPgAMaQ5MqFKH8m-XIybehSXgarDucSlmeK3E,37
+ftmq-4.3.4.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
+ftmq-4.3.4.dist-info/licenses/NOTICE,sha256=LNgfzuMbk3kIP_KnyDiXO8rQJmDrLy_PQ7cAY8lCmMM,463
+ftmq-4.3.4.dist-info/RECORD,,