Add static annotations for accumulators

third_party/3/pyspark/_typing.pyi

@@ -0,0 +1,6 @@
+from typing import Any, List, Optional, TypeVar, Union
+from typing_extensions import Protocol
+
+
+class SupportsIAdd(Protocol):
+ def __iadd__(self, other: SupportsIAdd) -> SupportsIAdd: ...

third_party/3/pyspark/accumulators.pyi

@@ -1,37 +1,45 @@
 # Stubs for pyspark.accumulators (Python 3.7)
 #
-# NOTE: This dynamically typed stub was automatically generated by stubgen.
+
+from typing import Callable, Generic, Tuple, Type, TypeVar
+
+import socketserver.BaseRequestHandler # type: ignore
+
+from pyspark._typing import SupportsIAdd
+
+T = TypeVar("T")
+U = TypeVar("U", bound=SupportsIAdd)
 
 import socketserver as SocketServer
 from typing import Any
 
-class Accumulator:
- aid: Any = ...
- accum_param: Any = ...
- def __init__(self, aid: Any, value: Any, accum_param: Any) -> None: ...
- def __reduce__(self): ...
+class Accumulator(Generic[T]):
+ aid = ... # type: int
+ accum_param = ... # type: AccumulatorParam[T]
+ def __init__(self, aid: int, value: T, accum_param: AccumulatorParam[T]) -> None: ...
+ def __reduce__(self) -> Tuple[Callable[[int, int, AccumulatorParam[T]], Accumulator[T]], Tuple[int, int, AccumulatorParam[T]]]: ...
  @property
- def value(self): ...
+ def value(self) -> T: ...
  @value.setter
- def value(self, value: Any) -> None: ...
- def add(self, term: Any) -> None: ...
- def __iadd__(self, term: Any): ...
+ def value(self, value: T) -> None: ...
+ def add(self, term: T) -> None: ...
+ def __iadd__(self, term: T) -> Accumulator[T]: ...
 
-class AccumulatorParam:
- def zero(self, value: Any) -> None: ...
- def addInPlace(self, value1: Any, value2: Any) -> None: ...
+class AccumulatorParam(Generic[T]):
+ def zero(self, value: T) -> T: ...
+ def addInPlace(self, value1: T, value2: T) -> T: ...
 
-class AddingAccumulatorParam(AccumulatorParam):
- zero_value: Any = ...
- def __init__(self, zero_value: Any) -> None: ...
- def zero(self, value: Any): ...
- def addInPlace(self, value1: Any, value2: Any): ...
+class AddingAccumulatorParam(AccumulatorParam[U]):
+ zero_value = ... # type: U
+ def __init__(self, zero_value: U) -> None: ...
+ def zero(self, value: U) -> U: ...
+ def addInPlace(self, value1: U, value2: U) -> U: ...
 
 class _UpdateRequestHandler(SocketServer.StreamRequestHandler):
- def handle(self): ...
+ def handle(self) -> None: ...
 
 class AccumulatorServer(SocketServer.TCPServer):
- auth_token: Any = ...
- def __init__(self, server_address: Any, RequestHandlerClass: Any, auth_token: Any) -> None: ...
+ auth_token = ... # type: str
+ def __init__(self, server_address: Tuple[str, int], RequestHandlerClass: Type[socketserver.BaseRequestHandler], auth_token: str) -> None: ...
  server_shutdown: bool = ...
  def shutdown(self) -> None: ...