Decimal-string display & casting layer for morton_index

mortie/morton_index.py

@@ -17,19 +17,95 @@
 clear ``ImportError`` is raised if the ExtensionArray is touched without pandas.
 """
 
+import os
+
 import numpy as np
 
 from . import _rustie
 
 # ``MortonIndexDtype`` / ``MortonIndexArray`` are provided via module-level
 # ``__getattr__`` (built lazily so a numpy-only install can import this module),
 # so they are intentionally not named in ``__all__`` here.
-__all__ = []
+__all__ = ["MortonIndexScalar"]
 
 # HEALPix orders this datatype reaches (0 = base cell, 29 = max resolution).
 MAX_ORDER = 29
 
 
+def _decimal_to_word(s):
+    """Parse a decimal Morton string back to its packed word (issue #104).
+
+    The inverse of the decode-through-kernel repr: sign + leading base digit
+    (``1..6``) then one ``1..4`` digit per order. Raises ``ValueError`` on any
+    malformed id. Order-29 strings parse to the *area* word -- the point/area
+    flag is not part of the decimal form, so the point word is unreachable
+    from its repr (the documented non-injectivity).
+    """
+    body = s[1:] if s.startswith("-") else s
+    if not (body.isdigit() and body.isascii()):
+        raise ValueError(f"malformed decimal Morton id {s!r}")
+    lead, digits = int(body[0]), body[1:]
+    if not 1 <= lead <= 6:
+        raise ValueError(
+            f"decimal Morton id {s!r}: base digit {lead} outside 1..6"
+        )
+    order = len(digits)
+    if order > MAX_ORDER:
+        raise ValueError(
+            f"decimal Morton id {s!r}: order {order} exceeds {MAX_ORDER}"
+        )
+    within = 0
+    for d in digits:
+        if not "1" <= d <= "4":
+            raise ValueError(
+                f"decimal Morton id {s!r}: digit {d} outside 1..4"
+            )
+        within = (within << 2) | (int(d) - 1)
+    base = lead + 5 if s.startswith("-") else lead - 1
+    nested = np.asarray([(base << (2 * order)) | within], dtype=np.uint64)
+    return int(_rustie.rust_mi_from_nested(nested, order)[0])
+
+
+class MortonIndexScalar(np.uint64):
+    """A packed ``morton_index`` word that displays as its decimal string.
+
+    Element access and iteration on a ``MortonIndexArray`` yield this type
+    (issue #104), so a downstream ``f"{shard_key}"`` prints the decimal Morton
+    id (``-31123`` style) rather than the raw packed word. It subclasses
+    ``numpy.uint64``: comparisons, hashing, and ``int()`` (the packed word)
+    behave exactly like the word itself; only ``str``/``repr`` differ. The
+    empty sentinel renders ``"<NA>"``; a word with an invalid prefix renders
+    ``"<invalid 0x...>"`` rather than raising (a repr must never raise).
+    """
+
+    def __str__(self):
+        word = int(self)
+        if word == 0:
+            return "<NA>"
+        try:
+            return _rustie.rust_mi_decimal_repr(
+                np.asarray([word], dtype=np.uint64)
+            )[0]
+        except ValueError:
+            return f"<invalid {word:#018x}>"
+
+    __repr__ = __str__
+
+    def __format__(self, spec):
+        # numpy's numeric __format__ would print the packed word; the display
+        # form of a morton_index is its decimal string, so f"{shard_key}" (and
+        # any string spec, e.g. ">10") formats that instead. int(self) remains
+        # the escape hatch to format the raw word numerically. Old-style
+        # "%d" % key bypasses __format__ entirely and emits the raw word.
+        return format(str(self), spec)
+
+    def __reduce__(self):
+        # numpy scalars pickle through multiarray.scalar, which rebuilds the
+        # bare np.uint64 and would silently drop the decimal display on any
+        # process boundary (multiprocessing/dask); rebuild the wrapper instead.
+        return (type(self), (int(self),))
+
+
 def _require_pandas():
     """Import pandas lazily, raising a clear error if it is absent.
 
@@ -272,7 +348,7 @@ def __len__(self):
         def __getitem__(self, item):
             result = self._data[item]
             if np.isscalar(result) or isinstance(result, np.integer):
-                return np.uint64(result)
+                return MortonIndexScalar(result)
             return type(self)(result)
 
         def __setitem__(self, key, value):
@@ -428,20 +504,130 @@ def decimal_repr(self):
             produced by *decoding* each word (the canonical render-only repr;
             backward-compatible with the legacy ``str(legacy_i64)`` for orders
             0..=18, the natural extension for 19..=29). Raises ``ValueError`` on
-            any empty / invalid word.
+            any empty / invalid word. Note the repr is not injective across
+            kinds: an order-29 *point* renders identically to the order-29
+            *area* on the same path (the point/area flag is not part of the
+            decimal form).
             """
             return _rustie.rust_mi_decimal_repr(self._data)
 
+        def to_decimal(self):
+            """Vectorized decimal-string emit as a fixed-width numpy array.
+
+            The always-strings interchange convention from issue #48: returns
+            the decode-through-kernel decimal strings as a ``"<U31"`` numpy
+            array (sign + base digit + 29 order digits is the widest form, so
+            the width is order-independent and stable across arrays). Raises
+            ``ValueError`` on any empty / invalid word. Same non-injectivity
+            note as :meth:`decimal_repr`: point and area render identically at
+            order 29.
+            """
+            return np.asarray(self.decimal_repr(), dtype="<U31")
+
+        def to_legacy_i64(self):
+            """Emit the legacy signed decimal ``int64`` form (orders <= 18).
+
+            The named escape hatch from issue #48's emit conventions --
+            interchange is always the decimal *string* (:meth:`to_decimal`),
+            storage the packed ``uint64``; this exists solely for testing new
+            output against old pinned values (pairing with
+            :meth:`from_legacy`). Any element above order 18 raises
+            ``ValueError`` (the legacy decimal overflows ``int64`` above
+            that), as does any empty / invalid word -- never truncated, never
+            data-dependent.
+            """
+            strings = self.decimal_repr()  # raises on empty / invalid words
+            orders = self.orders()
+            if len(self) and int(orders.max()) > 18:
+                raise ValueError(
+                    f"to_legacy_i64 is capped at order 18 (array holds order "
+                    f"{int(orders.max())}); use to_decimal() for orders 19-29"
+                )
+            return np.asarray([int(s) for s in strings], dtype=np.int64)
+
+        def hive_path(self, root="", suffix=".zarr"):
+            """Hive-layout path per element (issue #104; spec lands on #62).
+
+            The ``morton-hive/1`` convention (zagg's sparse-coverage design
+            record): one decimal digit per directory level, the full id as the
+            leaf inside its own node --
+            ``{root}/{sign+base}/{d1}/.../{d_order}/{full_id}{suffix}``, e.g.
+            ``-31123`` -> ``-3/1/1/2/3/-31123.zarr`` and the order-0 ``-3`` ->
+            ``-3/-3.zarr``. Every order is a node, so mixed shard orders nest
+            naturally: a coarser shard's leaf sits in the same directory its
+            finer siblings descend through. Returns a list of ``str``; raises
+            ``ValueError`` on any empty / invalid word.
+            """
+            prefix = root.rstrip("/") + "/" if root else ""
+            paths = []
+            for s in self.decimal_repr():
+                head = 2 if s.startswith("-") else 1
+                levels = "/".join([s[:head], *s[head:]])
+                paths.append(f"{prefix}{levels}/{s}{suffix}")
+            return paths
+
+        @classmethod
+        def from_hive_path(cls, paths, suffix=".zarr"):
+            """Parse hive-layout paths back to words (inverse of hive_path).
+
+            ``paths`` is a single slash-separated path (``str`` /
+            ``os.PathLike``) or an iterable of them. The leaf basename
+            carries the full decimal id; when the path also carries the digit
+            directories -- recognized by a ``{sign+base}``-shaped component
+            sitting at its slot above the leaf -- the whole chain is checked
+            against the leaf and a mis-filed leaf (wrong base cell, wrong
+            descent) raises ``ValueError``. A bare ``{full_id}.zarr``, or one
+            under an arbitrary root without the digit chain, skips the check.
+            Order-29 ids parse to the *area* word (see
+            :func:`_decimal_to_word`).
+            """
+            import pathlib
+
+            if isinstance(paths, (str, os.PathLike)):
+                paths = [paths]
+            words = []
+            for p in paths:
+                if isinstance(p, pathlib.PurePath):
+                    # honor the path's own flavor: a WindowsPath splits on
+                    # backslashes too, which a raw "/"-split would not
+                    parts = list(p.parts)
+                else:
+                    parts = str(p).rstrip("/").split("/")
+                leaf = parts[-1]
+                if suffix and not leaf.endswith(suffix):
+                    raise ValueError(
+                        f"hive leaf {leaf!r} does not end with {suffix!r}"
+                    )
+                dec = leaf[: len(leaf) - len(suffix)] if suffix else leaf
+                word = _decimal_to_word(dec)
+                head = 2 if dec.startswith("-") else 1
+                levels = [dec[:head], *dec[head:]]
+                # Enforce the directory cross-check only when the chain is
+                # anchored: a {sign+base}-shaped component (optional "-" plus
+                # one digit 1..6) at its expected slot. Anchoring on shape --
+                # not equality with the leaf's own sign+base -- keeps a
+                # mis-filed wrong-base leaf detectable while still treating
+                # an arbitrary root as skippable (root components are
+                # indistinguishable from digit dirs by count alone).
+                anchor = parts[-1 - len(levels)] if (
+                    len(parts) - 1 >= len(levels)
+                ) else ""
+                body = anchor[1:] if anchor.startswith("-") else anchor
+                if len(body) == 1 and "1" <= body <= "6":
+                    got = parts[-1 - len(levels):-1]
+                    if got != levels:
+                        raise ValueError(
+                            f"hive path {p!r} directories {'/'.join(got)!r} "
+                            f"do not match leaf id {dec!r}"
+                        )
+                words.append(word)
+            return cls(np.asarray(words, dtype=np.uint64))
+
         # -- repr ------------------------------------------------------------
 
         def _word_repr(self, word):
-            """Compact ``base/order`` label for one packed word."""
-            w = np.asarray([word], dtype=np.uint64)
-            if word == int(self._SENTINEL):
-                return "<NA>"
-            base = int(_rustie.rust_mi_base_cell_of(w)[0])
-            order = int(_rustie.rust_mi_order_of(w)[0])
-            return f"base={base} order={order}"
+            """Decimal-string label for one packed word (issue #104)."""
+            return str(MortonIndexScalar(word))
 
         def __repr__(self):
             n = len(self)