[BEAM-13439] Correct type signature on ptransform_fn to better support typecheckers

• Type: Improvement
• Status: Resolved
• Priority: P2
• Resolution: Fixed
• Affects Version/s: 2.34.0
• Fix Version/s: 2.37.0
• Component/s: sdk-py-core
• Labels: stale-P2 transform typecheck
• Assignee: Unassigned
• Reporter: Sandy Chapman
• Created: 10/Dec/21 13:38
• Updated: 15/Mar/22 18:49
• Resolved: 15/Mar/22 18:49

Description

Because the return type of the ptransform_fn is not specified, certain typecheckers will assume the type signature of functions decorated with it remain unchanged. However, ptransform_fn specifically adjusts the type signature such that the function no longer takes a pcoll.

This can be corrected by changing the the definition of ptransform_fn from:

def ptransform_fn(fn):

to:

def ptransform_fn(fn) -> Callable[..., _PTransformFnPTransform]: 

This allows the typechecker to understand the return type of a function decorated with ptransform_fn will instead be a _PTransformFnPTransform type.

Additionally, for Python 3.10, ParamSpec could be used to preserve the signature of the internal function as well.

from typing import ParamSpec, TypeVar, Concatenate

P = ParamSpec("P")
T = TypeVar("T")
R = TypeVar("R")

def ptransform_fn(fn: Callable[Concatenate[PCollection[T], P], R]) -> Callable[[P], _PTransformFnPTransform]: ...

Issue Links

• links to: GitHub Pull Request #16517