v0.3.4

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.3.3
+current_version = 0.3.4
 commit = True
 tag = True
 message = 🔖 Bump version: {current_version} → {new_version}

.pre-commit-config.yaml

@@ -5,7 +5,7 @@
 # * Run "pre-commit install".
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.4.0
+    rev: v4.5.0
     hooks:
       - id: check-added-large-files
       - id: check-case-conflict

.vscode/settings.json

@@ -0,0 +1,7 @@
+{
+    "python.testing.pytestArgs": [
+        "tests"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
+}

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 ## CHANGELOG
 
+### 0.3.4 (2023-10-16)
+
+- Add group based aggregation, aggregate by year
+- Add some util functions from downstream dependencies
+- Update dependencies
+
 ### 0.3.3 (2023-10-03)
 
 - Update dependencies

VERSION

@@ -1 +1 @@
-0.3.3
+0.3.4

ftmq/__init__.py

@@ -1,4 +1,4 @@
 from ftmq.query import Query
 
-__version__ = "0.3.3"
+__version__ = "0.3.4"
 __all__ = ["Query"]

ftmq/aggregations.py

@@ -1,18 +1,19 @@
 import statistics
 from collections import defaultdict
 from functools import cache
-from typing import Any, Iterable, TypeAlias
+from typing import Any, Generator, Iterable, TypeAlias
 
 from banal import ensure_list
 from followthemoney.schema import Schema
 from followthemoney.types import registry
 from pydantic import BaseModel
 
-from ftmq.enums import Aggregations, Properties
+from ftmq.enums import Aggregations, Fields, Properties
 from ftmq.types import CE, CEGenerator
-from ftmq.util import to_numeric
+from ftmq.util import clean_dict, to_numeric
 
 Value: TypeAlias = int | float | str
+Values: TypeAlias = list[Value]
 
 
 @cache
@@ -24,34 +25,60 @@ def get_is_numeric(schema: Schema, prop: str) -> bool:
 
 
 class Aggregation(BaseModel):
-    prop: Properties
+    prop: Properties | Fields
     func: Aggregations
-    values: list[int | float] = []
+    values: Values = []
     value: Value | None = None
+    group_props: list[Properties | Fields] | None = []
+    grouper: dict[Properties, dict[str, Values]] = defaultdict(
+        lambda: defaultdict(list)
+    )
+    groups: dict[Properties, dict[str, Value]] = defaultdict(dict)
 
     def __hash__(self) -> int:
-        return hash((self.prop, self.func))
+        return hash((self.prop, self.func, *sorted(ensure_list(self.group_props))))
 
     def __eq__(self, other: Any) -> bool:
         return hash(self) == hash(other)
 
-    def get_value(self) -> Value | None:
+    def get_value(self, values: Values) -> Value | None:
         if self.func == "min":
-            return min(self.values)
+            return min(values)
         if self.func == "max":
-            return max(self.values)
+            return max(values)
         if self.func == "sum":
-            return sum(self.values)
+            return sum(values)
         if self.func == "avg":
-            return statistics.mean(self.values)
+            return statistics.mean(values)
+        if self.func == "count":
+            return len(set(values))
+
+    def get_proxy_values(
+        self, proxy: CE, prop: Properties | Fields | None = None
+    ) -> Generator[str, None, None]:
+        prop = prop or self.prop
+        if prop == Fields.id:
+            yield proxy.id
+        elif prop == Fields.dataset:
+            yield from proxy.datasets
+        elif prop == Fields.schema:
+            yield proxy.schema.name
+        elif prop == Fields.year:
+            for value in proxy.get_type_values(registry.date):
+                yield value[:4]
+        else:
+            yield from proxy.get(prop, quiet=True)
 
     def collect(self, proxy: CE) -> CE:
         is_numeric = get_is_numeric(proxy.schema, self.prop)
-        for value in proxy.get(self.prop, quiet=True):
+        for value in self.get_proxy_values(proxy):
             if is_numeric:
                 value = to_numeric(value)
             if value is not None:
                 self.values.append(value)
+                for prop in self.group_props:
+                    for g in self.get_proxy_values(proxy, prop):
+                        self.grouper[prop][g].append(value)
         return proxy
 
     def apply(self, proxies: CEGenerator) -> CEGenerator:
@@ -63,14 +90,19 @@ def __enter__(self) -> "Aggregation":
         return self
 
     def __exit__(self, *args, **kwargs) -> None:
-        self.value = self.get_value()
+        self.value = self.get_value(self.values)
+        for prop in self.group_props:
+            for g, values in self.grouper[prop].items():
+                self.groups[prop][g] = self.get_value(values)
 
     def dict(self, *args, **kwargs) -> dict[str, Any]:
         self.__exit__()
         return super().dict(*args, **kwargs)
 
 
-AggregatorResult: TypeAlias = dict[Aggregations, dict[Properties, Value]]
+AggregatorResult: TypeAlias = dict[
+    Aggregations | dict[str, Aggregations], dict[Properties, Value]
+]
 
 
 class Aggregator(BaseModel):
@@ -81,8 +113,14 @@ def __enter__(self) -> "Aggregator":
         return self
 
     def __exit__(self, *args, **kwargs) -> None:
+        self.result["groups"] = defaultdict(lambda: defaultdict(dict))
         for agg in self.aggregations:
             self.result[str(agg.func)][str(agg.prop)] = agg.value
+            for group in agg.group_props:
+                self.result["groups"][str(group)][str(agg.func)][
+                    str(agg.prop)
+                ] = agg.groups[group]
+        self.result = clean_dict(self.result)
 
     def apply(self, proxies: CEGenerator) -> CEGenerator:
         for agg in self.aggregations:
@@ -91,17 +129,23 @@ def apply(self, proxies: CEGenerator) -> CEGenerator:
         self.__exit__()
 
     @classmethod
-    def from_dict(cls, data: dict[Aggregations, Iterable[Properties]]) -> "Aggregator":
+    def from_dict(
+        cls, data: dict[Aggregations | str, Iterable[Properties]]
+    ) -> "Aggregator":
+        groups = ensure_list(data.pop("groups", None))
         return cls(
             aggregations=[
-                Aggregation(prop=p, func=agg)
+                Aggregation(prop=p, func=agg, group_props=groups)
                 for agg, props in data.items()
                 for p in ensure_list(props)
-            ]
+            ],
         )
 
     def to_dict(self) -> dict[str, set[str]]:
         data = defaultdict(set)
+        data["groups"] = defaultdict(lambda: defaultdict(set))
         for agg in self.aggregations:
             data[str(agg.func)].add(str(agg.prop))
-        return dict(data)
+            for group in agg.group_props:
+                data["groups"][str(group)][str(agg.func)].add(str(agg.prop))
+        return clean_dict(data)