Skip to content

fix(python): correct type hint for to_tensor_fn parameter#5577

Merged
wjones127 merged 4 commits intolance-format:mainfrom
AndreaBozzo:fix/issue-3129-to-tensor-fn-type-hint
Jan 12, 2026
Merged

fix(python): correct type hint for to_tensor_fn parameter#5577
wjones127 merged 4 commits intolance-format:mainfrom
AndreaBozzo:fix/issue-3129-to-tensor-fn-type-hint

Conversation

@AndreaBozzo
Copy link
Copy Markdown
Contributor

@AndreaBozzo AndreaBozzo commented Dec 24, 2025

Fixes #3129

Fixes lance-format#3129

The type hint for  in  was
declaring a single-argument callable, but the function is actually
called with additional keyword arguments ( and
).

Changes:
- Update type hint to use  to allow arbitrary args
- Enhance docstring to document the expected function signature
@github-actions github-actions Bot added bug Something isn't working python labels Dec 24, 2025
Copy link
Copy Markdown
Contributor

@wjones127 wjones127 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we fix the type hint by making it correct and specific?

Comment thread python/python/lance/torch/data.py Outdated
Comment on lines +196 to +198
Callable[[pa.RecordBatch], Union[dict[str, torch.Tensor], torch.Tensor]]
Callable[
..., Union[dict[str, torch.Tensor], torch.Tensor]
]
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems wrong to just make this more generic. How about this?

Callable[[pa.RecordBatch, ...], Union[dict[str, torch.Tensor], torch.Tensor]]

Alternatively, can be fancier and do:

from typing import Protocol

class ToTensorFn(Protocol):
    def __call__(self, batch: pa.RecordBatch, hf_converter, use_blob_api: bool) -> Union[dict[str, torch.Tensor], torch.Tensor]: ...

...
to_tensor_fn: ToTensorFn

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for pointing this out, i agree about the hint being too generic.
I modified it with the ' less fancy' approach + ran lint in the last commit

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I thought about your comment and this is probably the cleanest way i can come up with, what do you think?

@wjones127 wjones127 self-assigned this Jan 1, 2026
@AndreaBozzo AndreaBozzo requested a review from wjones127 January 4, 2026 19:41
@wjones127 wjones127 merged commit 1ac488f into lance-format:main Jan 12, 2026
16 checks passed
@AndreaBozzo AndreaBozzo deleted the fix/issue-3129-to-tensor-fn-type-hint branch January 12, 2026 22:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working python

Projects

None yet

Development

Successfully merging this pull request may close these issues.

lance.torch.data.LanceDataset(to_tensor_fn=...) typehint inconsistent with usage.

2 participants