Feat: add support for operations on IRI lists, add pandas-like queries

setup.cfg

@@ -50,7 +50,7 @@ python_requires = >=3.10
 install_requires =
     importlib-metadata; python_version<"3.8"
     pydantic>=2.0.0
-    datamodel-code-generator>=0.51.0
+    datamodel-code-generator>=0.51.0,<0.55.0
     typing_extensions
     pyld
     rdflib

src/oold/backend/document_store.py

@@ -1,15 +1,27 @@
 import json
 import sqlite3
 from pathlib import Path
-from typing import Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Set, Union
 
-from oold.backend.interface import Backend, StoreResult
+from oold.backend.interface import (
+    Backend,
+    Condition,
+    LinkedDataFormat,
+    Query,
+    QueryParam,
+    ResolveParam,
+    ResolveResult,
+    StoreResult,
+    apply_operator,
+)
 
 
 class SimpleDictDocumentStore(Backend):
     _store: Optional[Dict[str, dict]] = None
+    format: LinkedDataFormat = LinkedDataFormat.JSON
 
-    def __init__(self):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
         self._store = {}
 
     def resolve_iris(self, iris: List[str]) -> Dict[str, Dict]:
@@ -18,13 +30,65 @@ def resolve_iris(self, iris: List[str]) -> Dict[str, Dict]:
             jsonld_dicts[iri] = self._store.get(iri, None)
         return jsonld_dicts
 
-    def store_jsonld_dicts(self, jsonld_dicts: Dict[str, Dict]) -> StoreResult:
-        for iri, jsonld_dict in jsonld_dicts.items():
-            self._store[iri] = jsonld_dict
+    def store_json_dicts(self, json_dicts: Dict[str, Dict]) -> StoreResult:
+        for iri, json_dict in json_dicts.items():
+            self._store[iri] = json_dict
         return StoreResult(success=True)
 
-    def query():
-        pass
+    def _filter(
+        self,
+        key: str,
+        operator: str,
+        value: Any,
+        context: Optional[Dict[str, Dict]] = None,
+        data: Optional[Dict[str, Dict]] = None,
+    ) -> Set[str]:
+        if data is None:
+            data = self._store
+        # retrieve property mapping from context
+        # ToDo: use a jsonld expand here
+        # if context is not None and key in context:
+        #    key = context[key]
+        matched_entities = set()
+        for iri, jsonld_dict in data.items():
+            if key in jsonld_dict:
+                if apply_operator(operator, jsonld_dict[key], value):
+                    matched_entities.add(iri)
+        return matched_entities
+
+    def _query(
+        self,
+        query: Union[Query, Condition],
+        context: Dict = None,
+        data: Optional[Dict[str, Dict]] = None,
+    ) -> Set[str]:
+        print("QUERY", query)
+        if data is None:
+            data = self._store
+        if isinstance(query, Condition):
+            return self._filter(query.field, query.operator, query.value, context, data)
+        elif isinstance(query, Query):
+            c1_res = self._query(query.op1, context, data)
+            c2_res = self._query(query.op2, context, data)
+            if query.operator == "and":
+                # intersect the results
+                return c1_res & c2_res
+            elif query.operator == "or":
+                # union the results
+                return c1_res | c2_res
+            else:
+                raise NotImplementedError(f"Operator {query.operator} not implemented")
+        else:
+            raise ValueError("Invalid query type")
+
+    def query(self, param: QueryParam) -> ResolveResult:
+        context = None
+        # if param.model_cls is not None:
+        #     context = _get_schema(param.model_cls).get("@context", None)
+        # elif self.model_cls is not None:
+        #     context = _get_schema(self.model_cls).get("@context", None)
+        iris = self._query(param.query, context)
+        return self.resolve(ResolveParam(iris=list(iris), model_cls=param.model_cls))
 
 
 class SqliteDocumentStore(Backend):

src/oold/backend/interface.py

@@ -1,11 +1,36 @@
+import operator as _op
 from abc import abstractmethod
+from enum import Enum
 from typing import Dict, List, Optional, Type, Union
 
 from pydantic import BaseModel
 
 from oold.static import GenericLinkedBaseModel
 
 
+class ComparisonOperator(str, Enum):
+    EQ = "eq"
+    NE = "ne"
+    LT = "lt"
+    LE = "le"
+    GT = "gt"
+    GE = "ge"
+
+
+_COMPARISON_FNS = {
+    ComparisonOperator.EQ: _op.eq,
+    ComparisonOperator.NE: _op.ne,
+    ComparisonOperator.LT: _op.lt,
+    ComparisonOperator.LE: _op.le,
+    ComparisonOperator.GT: _op.gt,
+    ComparisonOperator.GE: _op.ge,
+}
+
+
+def apply_operator(operator: ComparisonOperator, a, b) -> bool:
+    return _COMPARISON_FNS[operator](a, b)
+
+
 class SetResolverParam(BaseModel):
     iri: str
     resolver: "Resolver"
@@ -31,8 +56,45 @@ class ResolveResult(BaseModel):
     nodes: Dict[str, Union[None, GenericLinkedBaseModel]]
 
 
+class Query(BaseModel):
+    op1: Union["Query", "Condition"]
+    operator: str
+    op2: Union["Query", "Condition"]
+
+    # override the & operator
+    def __and__(self, other):
+        return Query(op1=self, operator="and", op2=other)
+
+
+class Condition(BaseModel):
+    field: str
+    operator: Optional[ComparisonOperator] = None
+    value: Optional[Union[str, int, float]] = None
+
+    # override the == operator
+    def __eq__(self, other):
+        self.operator = "eq"
+        self.value = other
+        return self
+
+    # override the & operator
+    def __and__(self, other):
+        return Query(op1=self, operator="and", op2=other)
+
+
+class QueryParam(BaseModel):
+    query: Union[Query, Condition]
+    model_cls: Optional[Type[GenericLinkedBaseModel]] = None
+
+
+class LinkedDataFormat(str, Enum):
+    JSON_LD = "JSON-LD"
+    JSON = "JSON"
+
+
 class Resolver(BaseModel):
     model_cls: Optional[Type[GenericLinkedBaseModel]] = None
+    format: Optional[LinkedDataFormat] = LinkedDataFormat.JSON_LD
 
     @abstractmethod
     def resolve_iris(self, iris: List[str]) -> Dict[str, Dict]:
@@ -53,11 +115,20 @@ def resolve(self, request: ResolveParam):
             if jsonld_dict is None:
                 nodes[iri] = None
             else:
-                node = model_cls.from_jsonld(jsonld_dict)
+                if self.format == LinkedDataFormat.JSON_LD:
+                    node = model_cls.from_jsonld(jsonld_dict)
+                elif self.format == LinkedDataFormat.JSON:
+                    node = model_cls.from_json(jsonld_dict)
+                else:
+                    raise ValueError(f"Unsupported format {self.format}")
                 nodes[iri] = node
 
         return ResolveResult(nodes=nodes)
 
+    def query(self, param: QueryParam) -> ResolveResult:
+        """Query the backend and return a ResolveResult."""
+        raise NotImplementedError("Query method not implemented in Resolver subclass")
+
 
 global _resolvers
 _resolvers = {}
@@ -100,35 +171,41 @@ class StoreResult(BaseModel):
     success: bool
 
 
-class Query(BaseModel):
-    pass
-
-
 class Backend(Resolver):
     def store(self, param: StoreParam) -> StoreResult:
         jsonld_dicts = {}
         for iri, node in param.nodes.items():
             if node is None:
                 jsonld_dicts[iri] = None
             else:
-                jsonld_dicts[iri] = node.to_jsonld()
-        return self.store_jsonld_dicts(jsonld_dicts)
+                if self.format == LinkedDataFormat.JSON_LD:
+                    jsonld_dicts[iri] = node.to_jsonld()
+                elif self.format == LinkedDataFormat.JSON:
+                    jsonld_dicts[iri] = node.to_json()
+                else:
+                    raise ValueError(f"Unsupported format {self.format}")
+        if self.format == LinkedDataFormat.JSON:
+            return self.store_json_dicts(jsonld_dicts)
+        else:
+            return self.store_jsonld_dicts(jsonld_dicts)
 
-    @abstractmethod
     def store_jsonld_dicts(self, jsonld_dicts: Dict[str, Dict]) -> StoreResult:
-        pass
+        raise NotImplementedError(
+            "store_jsonld_dicts method not implemented in Backend subclass"
+        )
 
-    @abstractmethod
-    def query(self, query: Query) -> ResolveResult:
-        """Query the backend and return a ResolveResult."""
-        pass
+    def store_json_dicts(self, json_dicts: Dict[str, Dict]) -> StoreResult:
+        raise NotImplementedError(
+            "store_json_dicts method not implemented in Backend subclass"
+        )
 
 
 global _backends
 _backends = {}
 
 
 def set_backend(param: SetBackendParam) -> None:
+    _resolvers[param.iri] = param.backend
     _backends[param.iri] = param.backend