Improve printing.char_from_number

Author: Armavica

pytensor/printing.py

@@ -63,23 +63,35 @@
 VALID_ASSOC = {"left", "right", "either"}
 
 
-def char_from_number(number):
-    """Convert numbers to strings by rendering it in base 26 using capital letters as digits."""
+def char_from_number(number: int) -> str:
+    """Convert a number to a string.
 
-    base = 26
+    It renders it in base 26 using capital letters as digits.
+    For example: 3·26² + 2·26¹ + 0·26⁰ → "DCA"
 
-    rval = ""
+    Parameters
+    ----------
+    number : int
+        The number to be converted.
 
-    if number == 0:
-        rval = "A"
+    Returns
+    -------
+    str
+        The converted string.
+    """
+
+    base = 26
+
+    remainders = []
 
     while number != 0:
-        remainder = number % base
-        new_char = chr(ord("A") + remainder)
-        rval = new_char + rval
-        number //= base
+        number, remainder = number // base, number % base
+        remainders.append(remainder)
 
-    return rval
+    if not remainders:
+        remainders = [0]
+
+    return "".join(chr(ord("A") + r) for r in remainders[::-1])
 
 
 @singledispatch

tests/test_printing.py

@@ -17,6 +17,7 @@
     PatternPrinter,
     PPrinter,
     Print,
+    char_from_number,
     debugprint,
     default_printer,
     get_node_by_id,
@@ -30,6 +31,22 @@
 from tests.graph.utils import MyInnerGraphOp, MyOp, MyVariable
 
 
+@pytest.mark.parametrize(
+    "number,s",
+    [
+        (0, "A"),
+        (1, "B"),
+        (25, "Z"),
+        (26, "BA"),
+        (27, "BB"),
+        (3 * 26**2 + 2 * 26 + 0, "DCA"),
+        (42421337, "DOVPLX"),
+    ],
+)
+def test_char_from_number(number: int, s: str):
+    assert char_from_number(number) == s
+
+
 @pytest.mark.skipif(not pydot_imported, reason="pydot not available")
 def test_pydotprint_cond_highlight():
     # This is a REALLY PARTIAL TEST.