fix metadata loss in to_geff

src/tracksdata/graph/_base_graph.py

@@ -1868,6 +1868,14 @@ def to_geff(
             It automatically generates the metadata with:
             - axes: time (t) and spatial axes ((z), y, x)
             - tracklet node property: tracklet_id
+            The graph metadata (`graph.metadata`) is always written to
+            `geff_metadata.extra["tracksdata"]`, including when the metadata is provided
+            by the caller. On key collisions the caller's value wins, so an explicit
+            `extra["tracksdata"]` entry still overrides the graph's metadata.
+            The caller's object is not modified.
+            `shape` is the canonical key for the shape of the dense segmentation, it is
+            read back by `GraphArrayView` and `to_ctc`. Use
+            `tracksdata.io.read_graph_metadata` to read it back without building a graph.
         overwrite : bool
             Whether to overwrite the geff data directory if it exists.
         zarr_format : Literal[2, 3]
@@ -1883,6 +1891,10 @@ def to_geff(
         edge_ids = edge_attrs.select(DEFAULT_ATTR_KEYS.EDGE_SOURCE, DEFAULT_ATTR_KEYS.EDGE_TARGET).to_numpy()
         edge_attrs = edge_attrs.drop(DEFAULT_ATTR_KEYS.EDGE_SOURCE, DEFAULT_ATTR_KEYS.EDGE_TARGET)
 
+        td_metadata = self.metadata.copy()
+        td_metadata.update(self._private_metadata_for_copy())
+        td_metadata.pop("geff", None)  # avoid geff being written multiple times
+
         if geff_metadata is None:
             axes = [Axis(name=DEFAULT_ATTR_KEYS.T, type="time")]
             axes.extend(
@@ -1917,10 +1929,6 @@ def to_geff(
                 for k, v in edge_attrs.to_dict().items()
             }
 
-            td_metadata = self.metadata.copy()
-            td_metadata.update(self._private_metadata_for_copy())
-            td_metadata.pop("geff", None)  # avoid geff being written multiple times
-
             geff_metadata = geff.GeffMetadata(
                 directed=True,
                 axes=axes,
@@ -1931,6 +1939,15 @@ def to_geff(
                     "tracksdata": td_metadata,
                 },
             )
+        else:
+            # copy so the caller's metadata object is left untouched
+            geff_metadata = geff_metadata.model_copy(deep=True)
+            extra = dict(geff_metadata.extra)
+            # caller-provided entries win, so they can still override the graph's metadata
+            merged = {**td_metadata, **extra.get("tracksdata", {})}
+            merged.pop("geff", None)  # avoid geff being written multiple times
+            extra["tracksdata"] = merged
+            geff_metadata.extra = extra
 
         node_dict = {
             k: {"values": column_to_numpy(v), "missing": None}

src/tracksdata/graph/_test/test_graph_backends.py

@@ -9,6 +9,7 @@
 import pytest
 import rustworkx as rx
 import sqlalchemy as sa
+from geff_spec import GeffMetadata
 from zarr.storage import MemoryStore
 
 from tracksdata.attrs import EdgeAttr, NodeAttr
@@ -2861,6 +2862,55 @@ def test_geff_roundtrip(graph_backend: BaseGraph) -> None:
     )
 
 
+def test_geff_roundtrip_custom_metadata(graph_backend: BaseGraph) -> None:
+    """Graph metadata must survive `to_geff` when the caller supplies its own `GeffMetadata`."""
+
+    _fill_mock_geff_graph(graph_backend)
+    graph_backend.metadata["shape"] = (5, 100, 100)
+
+    # a downstream library supplying its own metadata: same props as the auto-generated
+    # one, but with its own `extra` namespace instead of tracksdata's.
+    reference_store = MemoryStore()
+    graph_backend.to_geff(geff_store=reference_store)
+    custom_metadata = GeffMetadata.read(reference_store)
+    custom_metadata.extra = {"downstream": {"hello": "world"}}
+
+    output_store = MemoryStore()
+    graph_backend.to_geff(geff_store=output_store, geff_metadata=custom_metadata)
+
+    written_metadata = GeffMetadata.read(output_store)
+    # the caller's own namespace is untouched ...
+    assert written_metadata.extra["downstream"] == {"hello": "world"}
+    # ... and the graph metadata rode along. `shape` is a list, not a tuple, because
+    # the extras are serialized as JSON.
+    assert written_metadata.extra["tracksdata"]["shape"] == [5, 100, 100]
+
+    geff_graph, _ = IndexedRXGraph.from_geff(output_store)
+    assert geff_graph.metadata["shape"] == [5, 100, 100]
+
+    # the metadata object the caller passed in was not modified
+    assert custom_metadata.extra == {"downstream": {"hello": "world"}}
+
+
+def test_geff_custom_metadata_overrides_graph_metadata(graph_backend: BaseGraph) -> None:
+    """On key collisions the caller-supplied `extra["tracksdata"]` wins."""
+
+    _fill_mock_geff_graph(graph_backend)
+    graph_backend.metadata["shape"] = (5, 100, 100)
+
+    reference_store = MemoryStore()
+    graph_backend.to_geff(geff_store=reference_store)
+    custom_metadata = GeffMetadata.read(reference_store)
+    custom_metadata.extra = {"tracksdata": {"shape": [1, 2, 3], "extra_key": "value"}}
+
+    output_store = MemoryStore()
+    graph_backend.to_geff(geff_store=output_store, geff_metadata=custom_metadata)
+
+    geff_graph, _ = IndexedRXGraph.from_geff(output_store)
+    assert geff_graph.metadata["shape"] == [1, 2, 3]
+    assert geff_graph.metadata["extra_key"] == "value"
+
+
 def test_geff_overwrite(graph_backend: BaseGraph, tmp_path: Path) -> None:
     """Test that to_geff overwrites existing data in the store."""
     _fill_mock_geff_graph(graph_backend)

src/tracksdata/io/__init__.py

@@ -1,12 +1,17 @@
 """Input/output utilities for loading and saving tracking data in various formats."""
 
 from tracksdata.io._ctc import compressed_tracks_table, from_ctc, to_ctc
-from tracksdata.io._geff_dtypes import convert_geff_prop_dtype, geff_prop_dtype
+from tracksdata.io._geff import (
+    convert_geff_prop_dtype,
+    geff_prop_dtype,
+    read_graph_metadata,
+)
 
 __all__ = [
     "compressed_tracks_table",
     "convert_geff_prop_dtype",
     "from_ctc",
     "geff_prop_dtype",
+    "read_graph_metadata",
     "to_ctc",
 ]

src/tracksdata/io/_geff_dtypes.py → src/tracksdata/io/_geff.py

@@ -1,31 +1,79 @@
-"""Utilities for inspecting and converting the dtype of properties in geff files.
-
-The motivating case is segmentation masks: they are binary, but geff files
-written by older versions of tracksdata stored the mask ``data`` buffer as
-``uint64`` (see https://github.com/royerlab/tracksdata/pull/318). That is 8x
-larger than a boolean buffer both on disk and, more importantly, when read into
-memory, which can cause out-of-memory errors when loading large datasets. New
-files store masks as ``bool`` at write time, so :func:`convert_geff_prop_dtype`
-provides a one-time fix for legacy files.
-
-The helpers are not mask-specific: they read and rewrite the payload dtype of
-any geff property (node or edge, fixed- or variable-length). The caller names
-the property to act on.
+"""Utilities for inspecting and repairing geff datasets without loading the graph.
+
+:func:`read_graph_metadata` reads the tracksdata graph metadata of a geff dataset.
+`BaseGraph.to_geff` writes the graph metadata (`graph.metadata`) into the geff
+metadata extras and `BaseGraph.from_geff` hoists it back onto the graph, but callers
+that need a value *before* they have a graph object -- for example the ``shape`` of
+the dense segmentation, which is required to construct a `GraphArrayView` -- can use
+neither. This closes that gap so downstream libraries do not have to know where
+tracksdata stores the extras.
+
+:func:`geff_prop_dtype` and :func:`convert_geff_prop_dtype` inspect and convert the
+on-disk dtype of a property. The motivating case is segmentation masks: they are
+binary, but geff files written by older versions of tracksdata stored the mask
+``data`` buffer as ``uint64`` (see
+https://github.com/royerlab/tracksdata/pull/318). That is 8x larger than a boolean
+buffer both on disk and, more importantly, when read into memory, which can cause
+out-of-memory errors when loading large datasets. New files store masks as ``bool``
+at write time, so :func:`convert_geff_prop_dtype` provides a one-time fix for legacy
+files. The helpers are not mask-specific: they read and rewrite the payload dtype of
+any geff property (node or edge, fixed- or variable-length). The caller names the
+property to act on.
 """
 
 from __future__ import annotations
 
 import os
 import shutil
 from pathlib import Path
+from typing import Any
 
 import numpy as np
 import zarr
+from geff_spec import GeffMetadata
 from zarr.storage import StoreLike
 
+from tracksdata.graph._base_graph import BaseGraph
 from tracksdata.utils._logging import LOG
 
-__all__ = ["convert_geff_prop_dtype", "geff_prop_dtype"]
+__all__ = ["convert_geff_prop_dtype", "geff_prop_dtype", "read_graph_metadata"]
+
+_EXTRA_KEY = "tracksdata"
+
+
+def read_graph_metadata(source: StoreLike | GeffMetadata) -> dict[str, Any]:
+    """
+    Read the tracksdata graph metadata of a geff dataset without loading the graph.
+
+    Returns the same metadata that `graph.metadata` would hold after
+    `BaseGraph.from_geff`, minus the `geff` key. Note that the values went through a
+    JSON round-trip, so tuples come back as lists.
+
+    Parameters
+    ----------
+    source : StoreLike | GeffMetadata
+        The store or path of the geff dataset, or an already parsed `GeffMetadata`.
+
+    Returns
+    -------
+    dict[str, Any]
+        The graph metadata, empty if the dataset was not written by tracksdata.
+
+    Examples
+    --------
+    ```python
+    graph.metadata["shape"] = (5, 100, 100)
+    graph.to_geff("tracks.geff")
+
+    shape = read_graph_metadata("tracks.geff")["shape"]  # [5, 100, 100]
+    ```
+    """
+    if not isinstance(source, GeffMetadata):
+        source = GeffMetadata.read(source)
+
+    metadata = source.extra.get(_EXTRA_KEY, {})
+
+    return {k: v for k, v in metadata.items() if not BaseGraph._is_private_metadata_key(k)}
 
 
 def geff_prop_dtype(

src/tracksdata/io/_test/test_geff_dtypes.py

@@ -8,7 +8,7 @@
 from tracksdata.constants import DEFAULT_ATTR_KEYS
 from tracksdata.graph import IndexedRXGraph, RustWorkXGraph
 from tracksdata.io import convert_geff_prop_dtype, geff_prop_dtype
-from tracksdata.io._geff_dtypes import _overwrite_array, _set_prop_metadata_dtype
+from tracksdata.io._geff import _overwrite_array, _set_prop_metadata_dtype
 from tracksdata.nodes._mask import Mask
 
 MASK_KEY = DEFAULT_ATTR_KEYS.MASK

src/tracksdata/io/_test/test_geff_metadata.py

@@ -0,0 +1,89 @@
+from pathlib import Path
+
+import polars as pl
+from geff_spec import Axis, GeffMetadata, PropMetadata
+from zarr.storage import MemoryStore
+
+from tracksdata.graph import RustWorkXGraph
+from tracksdata.io import read_graph_metadata
+
+SHAPE = (5, 100, 100)
+
+
+def _make_graph() -> RustWorkXGraph:
+    graph = RustWorkXGraph()
+    graph.add_node_attr_key("y", pl.Float64)
+    graph.add_node_attr_key("x", pl.Float64)
+    graph.add_node({"t": 0, "y": 1.0, "x": 2.0})
+    graph.add_node({"t": 1, "y": 3.0, "x": 4.0})
+    graph.metadata["shape"] = SHAPE
+    return graph
+
+
+def _minimal_geff_metadata() -> GeffMetadata:
+    """A `GeffMetadata` as a downstream library would build it: no tracksdata extras."""
+    return GeffMetadata(
+        directed=True,
+        axes=[
+            Axis(name="t", type="time"),
+            Axis(name="y", type="space", scale=0.5),
+            Axis(name="x", type="space", scale=0.5),
+        ],
+        node_props_metadata={
+            "t": PropMetadata(identifier="t", dtype="int64"),
+            "y": PropMetadata(identifier="y", dtype="float64"),
+            "x": PropMetadata(identifier="x", dtype="float64"),
+        },
+        edge_props_metadata={},
+        extra={"downstream": {"hello": "world"}},
+    )
+
+
+def test_read_graph_metadata_from_path(tmp_path: Path) -> None:
+    """The shape is readable from a store path, without building a graph."""
+    graph = _make_graph()
+    geff_path = tmp_path / "tracks.geff"
+    graph.to_geff(geff_store=geff_path)
+
+    # tuples become lists through the JSON round-trip
+    assert read_graph_metadata(geff_path) == {"shape": list(SHAPE)}
+
+
+def test_read_graph_metadata_custom_geff_metadata() -> None:
+    """The shape survives a write with caller-supplied metadata and is readable back."""
+    graph = _make_graph()
+    store = MemoryStore()
+    graph.to_geff(geff_store=store, geff_metadata=_minimal_geff_metadata())
+
+    assert read_graph_metadata(store) == {"shape": list(SHAPE)}
+
+
+def test_read_graph_metadata_from_geff_metadata_instance() -> None:
+    """An already parsed `GeffMetadata` is accepted, so the store is not reopened."""
+    graph = _make_graph()
+    store = MemoryStore()
+    graph.to_geff(geff_store=store)
+
+    assert read_graph_metadata(GeffMetadata.read(store)) == {"shape": list(SHAPE)}
+
+
+def test_read_graph_metadata_without_tracksdata_extras() -> None:
+    """A geff not written by tracksdata yields an empty dict rather than raising."""
+    # foreign extras only
+    assert read_graph_metadata(_minimal_geff_metadata()) == {}
+
+    no_extra = _minimal_geff_metadata()
+    no_extra.extra = {}
+    assert read_graph_metadata(no_extra) == {}
+
+
+def test_read_graph_metadata_excludes_private_keys() -> None:
+    """Private metadata is written to the store but not exposed by the reader."""
+    graph = _make_graph()
+    graph._private_metadata["__private_secret"] = 42
+
+    store = MemoryStore()
+    graph.to_geff(geff_store=store)
+
+    assert "__private_secret" in GeffMetadata.read(store).extra["tracksdata"]
+    assert read_graph_metadata(store) == {"shape": list(SHAPE)}