Skip to content

Commit

Permalink
Revert typing changes to stats.py (#33622)
Browse files Browse the repository at this point in the history
* Revert typing changes to stats.py

* fix sequence hint
  • Loading branch information
jrmccluskey authored Jan 16, 2025
1 parent 7f9ab48 commit 6530d81
Showing 1 changed file with 13 additions and 11 deletions.
24 changes: 13 additions & 11 deletions sdks/python/apache_beam/transforms/stats.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,8 @@
import typing
from collections.abc import Callable
from typing import Any
from typing import List
from typing import Tuple

from apache_beam import coders
from apache_beam import typehints
Expand Down Expand Up @@ -319,8 +321,8 @@ def _display_data(num_quantiles, key, reverse, weighted, input_batched):
}

@typehints.with_input_types(
typehints.Union[typing.Sequence[T], typing.Tuple[T, float]])
@typehints.with_output_types(list[T])
typehints.Union[typing.Sequence[T], Tuple[T, float]])
@typehints.with_output_types(List[T])
class Globally(PTransform):
"""
PTransform takes PCollection and returns a list whose single value is
Expand Down Expand Up @@ -372,8 +374,8 @@ def display_data(self):
input_batched=self._input_batched)

@typehints.with_input_types(
typehints.Union[tuple[K, V], tuple[K, tuple[V, float]]])
@typehints.with_output_types(tuple[K, list[V]])
typehints.Union[Tuple[K, V], Tuple[K, Tuple[V, float]]])
@typehints.with_output_types(Tuple[K, List[V]])
class PerKey(PTransform):
"""
PTransform takes PCollection of KV and returns a list based on each key
Expand Down Expand Up @@ -450,7 +452,7 @@ def __init__(self, buffer_size, num_buffers, weighted, key, reverse):
self.less_than = lambda a, b: key(a) < key(b)

def get_argsort_key(self, elements):
# type: (list) -> Callable[[int], Any]
# type: (List) -> Callable[[int], Any]

"""Returns a key for sorting indices of elements by element's value."""
if self.key is None:
Expand Down Expand Up @@ -507,7 +509,7 @@ class _QuantileState(object):
Compact summarization of a collection on which quantiles can be estimated.
"""
def __init__(self, unbuffered_elements, unbuffered_weights, buffers, spec):
# type: (list, list, list[_QuantileBuffer], _QuantileSpec) -> None
# type: (List, List, List[_QuantileBuffer], _QuantileSpec) -> None
self.buffers = buffers
self.spec = spec
if spec.weighted:
Expand Down Expand Up @@ -541,7 +543,7 @@ def is_empty(self):
return not self.unbuffered_elements and not self.buffers

def _add_unbuffered(self, elements, offset_fn):
# type: (list, Any) -> None
# type: (List, Any) -> None

"""
Add elements to the unbuffered list, creating new buffers and
Expand All @@ -567,7 +569,7 @@ def _add_unbuffered(self, elements, offset_fn):
self.collapse_if_needed(offset_fn)

def _add_unbuffered_weighted(self, elements, offset_fn):
# type: (list, Any) -> None
# type: (List, Any) -> None

"""
Add elements with weights to the unbuffered list, creating new buffers and
Expand Down Expand Up @@ -654,7 +656,7 @@ def collapse_if_needed(self, offset_fn):


def _collapse(buffers, offset_fn, spec):
# type: (list[_QuantileBuffer], Any, _QuantileSpec) -> _QuantileBuffer
# type: (List[_QuantileBuffer], Any, _QuantileSpec) -> _QuantileBuffer

"""
Approximates elements from multiple buffers and produces a single buffer.
Expand Down Expand Up @@ -683,7 +685,7 @@ def _collapse(buffers, offset_fn, spec):


def _interpolate(buffers, count, step, offset, spec):
# type: (list[_QuantileBuffer], int, float, float, _QuantileSpec) -> tuple[list, list, Any, Any]
# type: (List[_QuantileBuffer], int, float, float, _QuantileSpec) -> Tuple[List, List, Any, Any]

"""
Emulates taking the ordered union of all elements in buffers, repeated
Expand Down Expand Up @@ -933,7 +935,7 @@ def add_input(self, quantile_state, element):
return quantile_state

def _add_inputs(self, quantile_state, elements):
# type: (_QuantileState, list) -> _QuantileState
# type: (_QuantileState, List) -> _QuantileState

"""
Add a batch of elements to the collection being summarized by quantile
Expand Down

0 comments on commit 6530d81

Please sign in to comment.