Start filling out set operations

docs/quickstart.md

@@ -9,7 +9,7 @@ jupytext:
 # Quickstart
 
 ## Installation
-```
+```bash
 pip install qubed
 ```
 
@@ -29,17 +29,49 @@ q = Qube.from_dict({
         "expver=0002": {"param=1":{}, "param=2":{}},
     },
 })
+print(f"{q.n_leaves = }, {q.n_nodes = }")
 q
 ```
 
-Compress the qube:
+Compress it:
 
 ```{code-cell} python3
-q.compress()
+cq = q.compress()
+assert cq.n_leaves == q.n_leaves
+print(f"{cq.n_leaves = }, {cq.n_nodes = }")
+cq
 ```
 
-Load some example qubes:
+Load a larger example qube (requires source checkout):
 
 ```{code-cell} python3
+from pathlib import Path
+import json
+data_path = Path("../tests/example_qubes/climate_dt.json")
+with data_path.open("r") as f:
+    climate_dt = Qube.from_json(json.loads(f.read()))
+
+# Using the html or print methods is optional but lets you specify things like the depth of the tree to display.
+print(f"{climate_dt.n_leaves = }, {climate_dt.n_nodes = }")
+climate_dt.html(depth=1) # Limit how much is open initially, click leave to see more.
+```
+
+### Set Operations
+
+```{code-cell} python3
+A = Qube.from_dict({
+    "a=1/2/3" : {"b=1/2/3" : {"c=1/2/3" : {}}},
+    "a=5" : {  "b=4" : {  "c=4" : {}}}
+    })
+
+B = Qube.from_dict({
+    "a=1/2/3" : {"b=1/2/3" : {"c=1/2/3" : {}}},
+    "a=5" : {  "b=4" : {  "c=4" : {}}}
+    })
+
+A.print(name="A"), B.print(name="B");
+
+A | B
+```
+
 
-### Set Operations

src/python/qubed/Qube.py

@@ -6,6 +6,7 @@ from typing import Any, Callable, Hashable, Literal, Mapping
 
 from frozendict import frozendict
 
+from . import set_operations
 from .tree_formatters import HTML, node_tree_to_html, node_tree_to_string
 from .value_types import DateRange, Enum, IntRange, TimeRange, Values
 
@@ -31,6 +32,12 @@ class NodeData:
 
     def summary(self) -> str:
         return f"{self.key}={self.values.summary()}" if self.key != "root" else "root"
+    
+@dataclass(frozen=True, eq=True, order=True)
+class RootNodeData(NodeData):
+    "Helper class to print a custom root name"
+    def summary(self) -> str:
+        return self.key
 
 @dataclass(frozen=True, eq=True, order=True)
 class Qube:
@@ -91,16 +98,21 @@ class Qube:
         return cls.make("root", Enum(("root",)), [])
 
     
-    def __str__(self, depth = None) -> str:
-        return "".join(node_tree_to_string(node=self, depth = depth))
+    def __str__(self, depth = None, name = None) -> str:
+        node = dataclasses.replace(self, data = RootNodeData(key = name, values=self.values, metadata=self.metadata)) if name is not None else self
+        return "".join(node_tree_to_string(node=node, depth = depth))
     
-    def print(self, depth = None): print(self.__str__(depth = depth))
+    def print(self, depth = None, name: str | None = None): 
+        print(self.__str__(depth = depth, name = name))
     
     def html(self, depth = 2, collapse = True) -> HTML:
         return HTML(node_tree_to_html(self, depth = depth, collapse = collapse))
     
     def _repr_html_(self) -> str:
         return node_tree_to_html(self, depth = 2, collapse = True)
+    
+    def __or__(self, other: "Qube") -> "Qube":
+        return set_operations.operation(self, other, set_operations.SetOperation.UNION)
 
     
     def __getitem__(self, args) -> 'Qube':
@@ -111,7 +123,13 @@ class Qube:
                 return dataclasses.replace(c, data = data)
         raise KeyError(f"Key {key} not found in children of {self.key}")
 
-    
+    @cached_property
+    def n_leaves(self) -> int:
+        return len(self.values) * (sum(c.n_leaves for c in self.children) if self.children else 1)
+
+    @cached_property
+    def n_nodes(self) -> int:
+        return 1 + sum(c.n_nodes for c in self.children)
 
     def transform(self, func: 'Callable[[Qube], Qube | list[Qube]]') -> 'Qube':
         """

src/python/qubed/set_operations.py

@@ -0,0 +1,21 @@
+from enum import Enum
+from collections import defaultdict
+
+
+class SetOperation(Enum):
+    UNION = (1, 1, 1)
+    INTERSECTION = (0, 1, 0)
+    DIFFERENCE = (1, 0, 0)
+    SYMMETRIC_DIFFERENCE = (1, 0, 1)
+
+
+def operation(A: "Qube", B : "Qube", type: SetOperation) -> "Qube":
+    # Sort nodes from both qubes by their keys
+    nodes_by_key = defaultdict(lambda : dict(A = [], B = []))
+    for node in A.nodes:
+        nodes_by_key[node.key]["A"].append(node)
+    for key, ndoes
+
+# The root node is special so we need a helper method that we can recurse on
+def _operation(A: list["Qube"], B : list["Qube"], type: SetOperation) -> "Qube":
+    pass

tests/test_basic_operations.py

@@ -0,0 +1,27 @@
+from qubed import Qube
+
+
+def test_eq():
+    d = {
+        "class=od" : {
+            "expver=0001": {"param=1":{}, "param=2":{}},
+            "expver=0002": {"param=1":{}, "param=2":{}},
+        },
+        "class=rd" : {
+            "expver=0001": {"param=1":{}, "param=2":{}, "param=3":{}},
+            "expver=0002": {"param=1":{}, "param=2":{}},
+        },
+    }
+    q = Qube.from_dict(d)
+    r = Qube.from_dict(d)
+
+    assert q == r
+
+def test_n_leaves():
+    q = Qube.from_dict({
+        "a=1/2/3" : {"b=1/2/3" : {"c=1/2/3" : {}}},
+        "a=5" : {  "b=4" : {  "c=4" : {}}}
+        })
+    
+    # Size is 3*3*3 + 1*1*1 = 27 + 1
+    assert q.n_leaves == 27 + 1