Add __repr__ and some pythonic interface

tiledb/consolidation_plan.py

@@ -1,3 +1,5 @@
+import pprint
+
 import tiledb.cc as lt
 
 from .ctx import Ctx, CtxMixin, default_ctx
@@ -25,6 +27,59 @@ def __init__(self, ctx: Ctx, array: lt.Array, fragment_size: int):
 
         super().__init__(ctx, lt.Array(ctx, array), fragment_size)
 
+    def __len__(self):
+        """Returns the number of nodes in the consolidation plan"""
+        return self.num_nodes
+
+    def __repr__(self):
+        attrs = {
+            "num_nodes": self.num_nodes,
+            "fragments": {
+                f"node_{node_idx}": {
+                    "num_fragments": self.num_fragments(node_idx),
+                    "fragment_uris": [
+                        self.fragment_uri(node_idx, fragment_idx)
+                        for fragment_idx in range(self.num_fragments(node_idx))
+                    ],
+                }
+                for node_idx in range(self.num_nodes)
+            },
+        }
+
+        return pprint.PrettyPrinter().pformat(attrs)
+
+    def _repr_html_(self):
+        from io import StringIO
+
+        output = StringIO()
+        output.write("<section>\n")
+
+        output.write("<h3>Consolidation Plan</h3>\n")
+        output.write("<table>\n")
+        output.write(
+            "<tr><th>Node</th><th>Num Fragments</th><th>Fragment URIs</th></tr>\n"
+        )
+        for node_idx in range(self.num_nodes):
+            output.write(
+                f"<tr><td>{node_idx}</td><td>{self.num_fragments(node_idx)}</td><td>{', '.join(self.fragment_uri(node_idx, fragment_idx) for fragment_idx in range(self.num_fragments(node_idx)))}</td></tr>\n"
+            )
+        output.write("</table>\n")
+
+        output.write("</section>\n")
+        return output.getvalue()
+
+    def __getitem__(self, idx):
+        if idx < 0 or idx >= self.num_nodes:
+            raise IndexError("Index out of bounds")
+
+        return {
+            "num_fragments": self.num_fragments(idx),
+            "fragment_uris": [
+                self.fragment_uri(idx, fragment_idx)
+                for fragment_idx in range(self.num_fragments(idx))
+            ],
+        }
+
     @property
     def num_nodes(self) -> int:
         """

tiledb/tests/test_consolidation_plan.py

@@ -1,6 +1,8 @@
 import json
+import xml
 
 import numpy as np
+import pytest
 
 import tiledb
 from tiledb.tests.common import DiskTestCase
@@ -16,25 +18,36 @@ def test_consolidation_plan(self):
         with tiledb.open(path, "r") as A:
             cons_plan = tiledb.ConsolidationPlan(tiledb.default_ctx(), A, 2)
             assert cons_plan.num_nodes == 1
+            assert cons_plan.num_nodes == len(cons_plan)
             assert cons_plan.num_fragments(0) == 1
             # check that it has a nodes key
             assert "nodes" in json.loads(cons_plan.dump())
             # check that it has a list of nodes
             assert isinstance(json.loads(cons_plan.dump())["nodes"], list)
             # check that each node has a uri key
-            assert "uri" in json.loads(cons_plan.dump())["nodes"][0]["uris"][0]
-
-        # write a second fragment to the array
+            for node in json.loads(cons_plan.dump())["nodes"]:
+                assert "uri" in node["uris"][0]
+            # test __repr__
+            try:
+                assert (
+                    xml.etree.ElementTree.fromstring(cons_plan._repr_html_())
+                    is not None
+                )
+            except:
+                pytest.fail(
+                    f"Could not parse cons_plan._repr_html_(). Saw {cons_plan._repr_html_()}"
+                )
+            # test __getitem__
+            assert cons_plan[0] == {
+                "num_fragments": 1,
+                "fragment_uris": [cons_plan.fragment_uri(0, 0)],
+            }
+
+        # write a second fragment to the array and check the new consolidation plan
         with tiledb.open(path, "w") as A:
             A[:] = np.random.rand(4)
 
         with tiledb.open(path, "r") as A:
             cons_plan = tiledb.ConsolidationPlan(tiledb.default_ctx(), A, 4)
             assert cons_plan.num_nodes == 1
             assert cons_plan.num_fragments(0) == 2
-            # check that it has a nodes key
-            assert "nodes" in json.loads(cons_plan.dump())
-            # check that it has a list of nodes
-            assert isinstance(json.loads(cons_plan.dump())["nodes"], list)
-            # check that each node has a uri key
-            assert "uri" in json.loads(cons_plan.dump())["nodes"][0]["uris"][0]