Merge pull request #560 from jdegenstein/groupby_str_repr

Better GroupBy dunder str/repr/ipython methods

src/build123d/topology.py

@@ -64,6 +64,7 @@
 from typing_extensions import Self, Literal
 
 from anytree import NodeMixin, PreOrderIter, RenderTree
+from IPython.lib.pretty import pretty
 from scipy.spatial import ConvexHull
 from vtkmodules.vtkCommonDataModel import vtkPolyData
 from vtkmodules.vtkFiltersCore import vtkPolyDataNormals, vtkTriangleFilter
@@ -3750,6 +3751,23 @@ def __len__(self):
     def __getitem__(self, key: int):
         return self.groups[key]
 
+    def __str__(self):
+        return pretty(self)
+
+    def __repr__(self):
+        return repr(ShapeList(self))
+
+    def _repr_pretty_(self, p, cycle = False):
+        if cycle:
+            p.text('(...)')
+        else:
+            with p.group(1, '[', ']'):
+                for idx, item in enumerate(self):
+                    if idx:
+                        p.text(',')
+                        p.breakable()
+                    p.pretty(item)
+
     def group(self, key: K):
         """Select group by key"""
         for k, i in self.key_to_group_index:

tests/test_direct_api.py

@@ -1,5 +1,6 @@
 # system modules
 import copy
+import io
 import json
 import math
 import os
@@ -9,6 +10,7 @@
 from typing import Optional
 import unittest
 from random import uniform
+from IPython.lib import pretty
 
 from OCP.BRepBuilderAPI import BRepBuilderAPI_MakeEdge
 from OCP.gp import (
@@ -2853,6 +2855,21 @@ def test_manifold(self):
 class TestShapeList(DirectApiTestCase):
     """Test ShapeList functionality"""
 
+    def assertDunderStrEqual(self, actual: str, expected_lines: list[str]):
+        actual_lines = actual.splitlines()
+        self.assertEqual(len(actual_lines), len(expected_lines))
+        for actual_line, expected_line in zip(actual_lines, expected_lines):
+            start, end = re.split(r"at 0x[0-9a-f]+", expected_line, 2, re.I)
+            self.assertTrue(actual_line.startswith(start))
+            self.assertTrue(actual_line.endswith(end))
+
+    def assertDunderReprEqual(self, actual: str, expected: str):
+        splitter = r"at 0x[0-9a-f]+"
+        actual_split_list = re.split(splitter, actual, 0, re.I)
+        expected_split_list = re.split(splitter, expected, 0, re.I)
+        for actual_split, expected_split in zip(actual_split_list, expected_split_list):
+            self.assertEqual(actual_split, expected_split)
+
     def test_sort_by(self):
         faces = Solid.make_box(1, 2, 3).faces() < SortBy.AREA
         self.assertAlmostEqual(faces[-1].area, 2, 5)
@@ -2963,6 +2980,41 @@ def test_group_by_retrieve_groups(self):
         with self.assertRaises(KeyError):
             result.group("C")
 
+    def test_group_by_str_repr(self):
+        nonagon = RegularPolygon(5,9)
+
+        expected = [
+            "[[<build123d.topology.Edge at 0x1277f6e1cd0>],",
+            " [<build123d.topology.Edge at 0x1277f6e1c10>,",
+            "  <build123d.topology.Edge at 0x1277fd8a090>],",
+            " [<build123d.topology.Edge at 0x1277f75d690>,",
+            "  <build123d.topology.Edge at 0x127760d9310>],",
+            " [<build123d.topology.Edge at 0x12777261f90>,",
+            "  <build123d.topology.Edge at 0x1277f6bd2d0>],",
+            " [<build123d.topology.Edge at 0x1276fbb0590>,",
+            "  <build123d.topology.Edge at 0x1277fec6d90>]]",
+        ]
+
+        self.assertDunderStrEqual(str(nonagon.edges().group_by(Axis.X)), expected)
+
+        expected_repr = (
+            "[[<build123d.topology.Edge object at 0x000001277FEC6D90>],"
+            " [<build123d.topology.Edge object at 0x000001277F6BCC10>,"
+            " <build123d.topology.Edge object at 0x000001277EC3D5D0>],"
+            " [<build123d.topology.Edge object at 0x000001277F6BEA90>,"
+            " <build123d.topology.Edge object at 0x000001276FCB2310>],"
+            " [<build123d.topology.Edge object at 0x000001277F6D10D0>,"
+            " <build123d.topology.Edge object at 0x000001276FBAAD10>],"
+            " [<build123d.topology.Edge object at 0x000001277FC86F90>,"
+            " <build123d.topology.Edge object at 0x000001277F6E1CD0>]]"
+        )
+        self.assertDunderReprEqual(repr(nonagon.edges().group_by(Axis.X)),expected_repr)
+
+        f = io.StringIO()
+        p = pretty.PrettyPrinter(f)
+        nonagon.edges().group_by(Axis.X)._repr_pretty_(p, cycle=True)
+        self.assertEqual(f.getvalue(), "(...)")
+
     def test_distance(self):
         with BuildPart() as box:
             Box(1, 2, 3)