fix(python): keep legacy _zvec pickle paths loadable after the move

pybind11 records a class's defining module inside the pickle stream, so objects
pickled by zvec <= 0.5.x reference the old top-level extension paths
(`_zvec.param` / `_zvec.schema` / `_zvec.typing`). After moving the extension to
`zvec._zvec`, loading such a pickle failed (ModuleNotFoundError, or a pybind
"type already registered" error from re-importing the bare `_zvec`).

Register `sys.modules` aliases for the legacy names in `zvec/__init__.py` so old
pickles still load after upgrading, and add a regression test covering the
old-layout -> new-layout pickle path.

Author: Cuiyus

python/tests/test_pickle_compat.py

@@ -0,0 +1,48 @@
+# Copyright 2025-present the zvec project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Legacy-pickle compatibility after the compiled extension moved from the
+top-level ``_zvec`` module to ``zvec._zvec``.
+
+pybind11 records a class's defining module inside the pickle stream, so objects
+pickled by zvec <= 0.5.x reference ``_zvec.param`` / ``_zvec.schema`` /
+``_zvec.typing``. ``zvec/__init__.py`` registers ``sys.modules`` aliases so
+those pickles still load after upgrading.
+"""
+
+from __future__ import annotations
+
+import pickle
+import sys
+
+from zvec.model.param import HnswIndexParam
+
+
+class TestLegacyPickleCompat:
+    def test_legacy_module_aliases_registered(self):
+        ext = sys.modules["zvec._zvec"]
+        assert sys.modules.get("_zvec") is ext
+        for sub in ("param", "schema", "typing"):
+            assert sys.modules.get(f"_zvec.{sub}") is getattr(ext, sub)
+
+    def test_old_layout_pickle_still_loads(self):
+        param = HnswIndexParam(m=16, ef_construction=200)
+        # protocol 2 encodes globals as ``c<module>\n<qualname>\n`` (no length
+        # prefix / framing), so we can rewrite the module path to the pre-0.5
+        # top-level form to forge what an old wheel would have produced.
+        blob = pickle.dumps(param, protocol=2)
+        assert b"czvec._zvec.param\n" in blob
+        legacy = blob.replace(b"czvec._zvec.param\n", b"c_zvec.param\n")
+
+        restored = pickle.loads(legacy)
+        assert restored.m == 16

python/zvec/__init__.py

@@ -58,6 +58,21 @@
     load_diskann_plugin,
 )
 
+# —— Backwards compatibility for legacy pickles ——
+# zvec <= 0.5.x shipped the compiled extension at the top level as ``_zvec``
+# (with submodules ``_zvec.param`` / ``_zvec.schema`` / ``_zvec.typing``); it
+# now lives at ``zvec._zvec``. pybind11 records a class's defining module inside
+# the pickle stream, so objects pickled by an old wheel reference the legacy
+# paths. Alias the old names in ``sys.modules`` so those pickles still load
+# after upgrading. (``zvec._zvec`` is already imported by the re-exports above.)
+_legacy_ext = sys.modules["zvec._zvec"]
+sys.modules.setdefault("_zvec", _legacy_ext)
+for _legacy_sub in ("param", "schema", "typing"):
+    _legacy_mod = getattr(_legacy_ext, _legacy_sub, None)
+    if _legacy_mod is not None:
+        sys.modules.setdefault(f"_zvec.{_legacy_sub}", _legacy_mod)
+del _legacy_ext, _legacy_sub, _legacy_mod
+
 from . import model as model
 
 # —— Extensions ——