fix(python): bind packaged BlobReader annotations to the worker type (#4726)

Packaged UDFs omit globals used only in postponed annotations, so a
`BlobReader` input has a correct wire signature but
`typing.get_type_hints()` on the loaded function raises `NameError`.

Use the resolved input signature to bind reader annotations to an
imported `lance_blob.BlobReader` alias. Preserve signature nullability
and avoid collisions with existing globals. Regression tests cover
functions and callable classes, eager and postponed annotations,
aliases, forward references, qualified names, and body globals.

Validation:

- Confirmed the new regression fails with `NameError: name 'BlobReader'
is not defined` against the original packager at
`057a39a1c70e15c78461bb7e3db8cef4c067cd3e`.
- The packaging, first-class Function slice 1 and 2, and BlobReader
suites pass all 159 tests with Python file-import directory caches
refreshed before lookups. The standard run passes 156 tests; three
temporary-module import failures reproduce unchanged on the original
packager because directory timestamps do not advance between file writes
in this environment. All 37 new regression cases pass in the standard
run.
- Python Ruff lint and formatting checks pass. Repository-wide Ruff lint
reports 48 existing errors in documentation examples.

Fixes #4718

<!-- lance-gatekeeper-fix:v1 agent=a32bb6efa79217b3c1c5c9d7a7ff9928
generation=1 -->

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>

python/python/lancedb/functions.py

@@ -1462,8 +1462,8 @@ class _SourcePackager:
         self._packaged: set[int] = set()
         self._names: dict[str, Any] = {}
 
-    def package(self, target: Any) -> str:
-        entry = self._definition(target, entry=True)
+    def package(self, target: Any, signature: FunctionSignature) -> str:
+        entry = self._definition(target, entry=True, signature=signature)
         parts = [_SOURCE_HEADER]
         if self._imports:
             parts.extend(["", *(self._imports[name] for name in sorted(self._imports))])
@@ -1472,12 +1472,20 @@ class _SourcePackager:
         parts.extend(["", entry, ""])
         return "\n".join(parts)
 
-    def _definition(self, target: Any, *, entry: bool) -> str:
+    def _definition(
+        self, target: Any, *, entry: bool, signature: Optional[FunctionSignature] = None
+    ) -> str:
         definition = _definition_node(target, entry=entry)
         if not inspect.isclass(target) and inspect.getclosurevars(target).nonlocals:
             raise ValueError(
                 "@udf cannot package functions that capture closure values"
             )
+        namespace = _definition_namespace(target)
+        reader_name = (
+            self._blob_annotations(definition, signature, namespace)
+            if signature is not None
+            else None
+        )
         source = ast.unparse(definition)
         references = _module_references(f"{_SOURCE_HEADER}\n\n{source}\n")
         dynamic = _namespace_acquisition(definition, references)
@@ -1485,9 +1493,10 @@ class _SourcePackager:
             raise ValueError(
                 f"@udf cannot package dynamic namespace access: {dynamic!r}"
             )
-        namespace = _definition_namespace(target)
         self._packaged.add(id(target))
         self._claim(target.__name__, ("definition", id(target)))
+        if reader_name is not None:
+            self._bind(reader_name, BlobReader)
         unresolved = []
         for name in sorted(references):
             if name == target.__name__:
@@ -1508,6 +1517,44 @@ class _SourcePackager:
             )
         return source
 
+    def _blob_annotations(
+        self,
+        definition: Union[ast.FunctionDef, ast.ClassDef],
+        signature: FunctionSignature,
+        namespace: Mapping[str, Any],
+    ) -> Optional[str]:
+        """Bind reader inputs to the worker type, including postponed aliases."""
+        readers = {
+            parameter.name: parameter
+            for parameter in signature.inputs
+            if parameter.blob_access == "reader"
+        }
+        if not readers:
+            return None
+        call = definition
+        if isinstance(definition, ast.ClassDef):
+            call = next(
+                (
+                    member
+                    for member in definition.body
+                    if isinstance(member, ast.FunctionDef) and member.name == "__call__"
+                ),
+                None,
+            )
+        if call is None:
+            return None
+        # Avoid shadowing globals used by the body or other packaged definitions.
+        name = "_lance_blob_reader"
+        while name in namespace or name in self._names or name == definition.name:
+            name += "_"
+        for argument in (*call.args.args, *call.args.kwonlyargs):
+            if argument.arg in readers:
+                annotation = name
+                if readers[argument.arg].nullable:
+                    annotation += " | None"
+                argument.annotation = ast.parse(annotation, mode="eval").body
+        return name
+
     def _claim(self, name: str, meaning: Any) -> None:
         claimed = self._names.setdefault(name, meaning)
         if claimed != meaning:
@@ -1845,7 +1892,7 @@ class UdfDefinition:
         code_files, code_modules = _code_files(code)
         entry_source = _SourcePackager(
             code_modules, frozenset(_requirement_name(package) for package in packages)
-        ).package(function)
+        ).package(function, signature)
         if code_files:
             kind = _PYTHON_BUNDLE_ARTIFACT
             source = json.dumps(

python/python/tests/test_function_packaging.py

@@ -11,9 +11,10 @@ import json
 import subprocess
 import sys
 import textwrap
+import types
 import warnings
 from pathlib import Path
-from typing import Annotated, Optional
+from typing import Annotated, Optional, get_type_hints
 
 import pyarrow as pa
 import pytest
@@ -126,6 +127,95 @@ def test_function_without_initialization_keeps_its_wire_form():
     assert request["artifact"]["kind"] == "python_callable"
 
 
+@pytest.mark.parametrize("class_udf", [False, True])
+@pytest.mark.parametrize("postponed", [False, True])
+@pytest.mark.parametrize(
+    ("annotation", "nullable"),
+    [
+        ("BlobReader", False),
+        ("Reader", False),
+        ('"Reader"', False),
+        ("Reader | None", True),
+        ("Optional[Reader]", True),
+        ('Optional["Reader"]', True),
+        ("MaybeReader", True),
+        ("lancedb.BlobReader", False),
+        ("lancedb.BlobReader | None", True),
+    ],
+)
+def test_blob_reader_annotations_resolve_in_the_artifact(
+    helper_package, monkeypatch, class_udf, postponed, annotation, nullable
+):
+    definition = (
+        f"class read_blob:\n"
+        f"    def __call__(self, video: {annotation}) -> int:\n"
+        f"        return 0\n"
+        if class_udf
+        else f"def read_blob(video: {annotation}) -> int:\n    return 0\n"
+    )
+    module = helper_package(
+        "fx_blob_annotations",
+        {
+            "__init__.py": (
+                ("from __future__ import annotations\n" if postponed else "")
+                + "import lancedb\n"
+                "from typing import Optional\n"
+                "from lancedb import BlobReader, BlobReader as Reader\n"
+                "from lancedb.functions import udf\n"
+                "MaybeReader = Optional[Reader]\n\n"
+                f"@udf\n{definition}"
+            )
+        },
+    )
+    request = module.read_blob.registration_request
+    parameter = request.signature.inputs[0]
+    assert parameter.arrow_type == "blob_v2"
+    assert parameter.blob_access == "reader"
+    assert parameter.nullable is nullable
+
+    # Workers provide lance_blob rather than importing the client SDK's reader.
+    worker = types.ModuleType("lance_blob")
+    worker.BlobReader = type("BlobReader", (), {})
+    monkeypatch.setitem(sys.modules, "lance_blob", worker)
+    namespace = {}
+    exec(compile(artifact(module.read_blob), "user_function.py", "exec"), namespace)
+    packaged = namespace["read_blob"]
+    if class_udf:
+        packaged = packaged.__call__
+    assert get_type_hints(packaged) == {
+        "video": Optional[worker.BlobReader] if nullable else worker.BlobReader,
+        "return": int,
+    }
+
+
+def test_blob_reader_annotation_does_not_shadow_body_globals(
+    helper_package, monkeypatch
+):
+    module = helper_package(
+        "fx_blob_annotation_globals",
+        {
+            "__init__.py": """
+                from lancedb import BlobReader
+                from lancedb.functions import udf
+
+                _lance_blob_reader = 7
+
+                @udf
+                def read_blob(video: BlobReader) -> int:
+                    return _lance_blob_reader
+            """
+        },
+    )
+    worker = types.ModuleType("lance_blob")
+    worker.BlobReader = type("BlobReader", (), {})
+    monkeypatch.setitem(sys.modules, "lance_blob", worker)
+    namespace = {}
+    exec(compile(artifact(module.read_blob), "user_function.py", "exec"), namespace)
+    packaged = namespace["read_blob"]
+    assert get_type_hints(packaged)["video"] is worker.BlobReader
+    assert packaged(worker.BlobReader()) == 7
+
+
 @pytest.mark.parametrize(
     ("init", "message"),
     [