Merge pull request #31755 Modernize type hints.

sdks/python/apache_beam/coders/observable_test.py

@@ -29,7 +29,7 @@
 class ObservableMixinTest(unittest.TestCase):
  observed_count = 0
  observed_sum = 0
- observed_keys = [] # type: List[Optional[str]]
+ observed_keys: List[Optional[str]] = []
 
  def observer(self, value, key=None):
  self.observed_count += 1

sdks/python/apache_beam/coders/row_coder.py

@@ -117,8 +117,7 @@ def from_type_hint(cls, type_hint, registry):
  return cls(schema)
 
  @staticmethod
- def from_payload(payload):
- # type: (bytes) -> RowCoder
+ def from_payload(payload: bytes) -> 'RowCoder':
  return RowCoder(proto_utils.parse_Bytes(payload, schema_pb2.Schema))
 
  def __reduce__(self):

sdks/python/apache_beam/coders/slow_stream.py

@@ -30,11 +30,10 @@ class OutputStream(object):
 
  A pure Python implementation of stream.OutputStream."""
  def __init__(self):
- self.data = [] # type: List[bytes]
+ self.data: List[bytes] = []
  self.byte_count = 0
 
- def write(self, b, nested=False):
- # type: (bytes, bool) -> None
+ def write(self, b: bytes, nested: bool = False) -> None:
  assert isinstance(b, bytes)
  if nested:
  self.write_var_int64(len(b))
@@ -45,8 +44,7 @@ def write_byte(self, val):
  self.data.append(chr(val).encode('latin-1'))
  self.byte_count += 1
 
- def write_var_int64(self, v):
- # type: (int) -> None
+ def write_var_int64(self, v: int) -> None:
  if v < 0:
  v += 1 << 64
  if v <= 0:
@@ -78,16 +76,13 @@ def write_bigendian_double(self, v):
  def write_bigendian_float(self, v):
  self.write(struct.pack('>f', v))
 
- def get(self):
- # type: () -> bytes
+ def get(self) -> bytes:
  return b''.join(self.data)
 
- def size(self):
- # type: () -> int
+ def size(self) -> int:
  return self.byte_count
 
- def _clear(self):
- # type: () -> None
+ def _clear(self) -> None:
  self.data = []
  self.byte_count = 0
 
@@ -101,8 +96,7 @@ def __init__(self):
  super().__init__()
  self.count = 0
 
- def write(self, byte_array, nested=False):
- # type: (bytes, bool) -> None
+ def write(self, byte_array: bytes, nested: bool = False) -> None:
  blen = len(byte_array)
  if nested:
  self.write_var_int64(blen)
@@ -125,25 +119,21 @@ class InputStream(object):
  """For internal use only; no backwards-compatibility guarantees.
 
  A pure Python implementation of stream.InputStream."""
- def __init__(self, data):
- # type: (bytes) -> None
+ def __init__(self, data: bytes) -> None:
  self.data = data
  self.pos = 0
 
  def size(self):
  return len(self.data) - self.pos
 
- def read(self, size):
- # type: (int) -> bytes
+ def read(self, size: int) -> bytes:
  self.pos += size
  return self.data[self.pos - size:self.pos]
 
- def read_all(self, nested):
- # type: (bool) -> bytes
+ def read_all(self, nested: bool) -> bytes:
  return self.read(self.read_var_int64() if nested else self.size())
 
- def read_byte(self):
- # type: () -> int
+ def read_byte(self) -> int:
  self.pos += 1
  return self.data[self.pos - 1]
 

sdks/python/apache_beam/coders/standard_coders_test.py

@@ -300,7 +300,7 @@ def json_value_parser(self, coder_spec):
  # Used when --fix is passed.
 
  fix = False
- to_fix = {} # type: Dict[Tuple[int, bytes], bytes]
+ to_fix: Dict[Tuple[int, bytes], bytes] = {}
 
  @classmethod
  def tearDownClass(cls):

sdks/python/apache_beam/coders/typecoders.py

@@ -80,8 +80,8 @@ def MakeXyzs(v):
 class CoderRegistry(object):
  """A coder registry for typehint/coder associations."""
  def __init__(self, fallback_coder=None):
- self._coders = {} # type: Dict[Any, Type[coders.Coder]]
- self.custom_types = [] # type: List[Any]
+ self._coders: Dict[Any, Type[coders.Coder]] = {}
+ self.custom_types: List[Any] = []
  self.register_standard_coders(fallback_coder)
 
  def register_standard_coders(self, fallback_coder):
@@ -104,12 +104,14 @@ def register_standard_coders(self, fallback_coder):
  def register_fallback_coder(self, fallback_coder):
  self._fallback_coder = FirstOf([fallback_coder, self._fallback_coder])
 
- def _register_coder_internal(self, typehint_type, typehint_coder_class):
- # type: (Any, Type[coders.Coder]) -> None
+ def _register_coder_internal(
+ self, typehint_type: Any,
+ typehint_coder_class: Type[coders.Coder]) -> None:
  self._coders[typehint_type] = typehint_coder_class
 
- def register_coder(self, typehint_type, typehint_coder_class):
- # type: (Any, Type[coders.Coder]) -> None
+ def register_coder(
+ self, typehint_type: Any,
+ typehint_coder_class: Type[coders.Coder]) -> None:
  if not isinstance(typehint_coder_class, type):
  raise TypeError(
  'Coder registration requires a coder class object. '
@@ -122,8 +124,7 @@ def register_coder(self, typehint_type, typehint_coder_class):
  typehint_type = getattr(typehint_type, '__name__', str(typehint_type))
  self._register_coder_internal(typehint_type, typehint_coder_class)
 
- def get_coder(self, typehint):
- # type: (Any) -> coders.Coder
+ def get_coder(self, typehint: Any) -> coders.Coder:
  if typehint and typehint.__module__ == '__main__':
  # See https://github.com/apache/beam/issues/21541
  # TODO(robertwb): Remove once all runners are portable.
@@ -187,8 +188,7 @@ class FirstOf(object):
  """For internal use only; no backwards-compatibility guarantees.
 
  A class used to get the first matching coder from a list of coders."""
- def __init__(self, coders):
- # type: (Iterable[Type[coders.Coder]]) -> None
+ def __init__(self, coders: Iterable[Type[coders.Coder]]) -> None:
  self._coders = coders
 
  def from_type_hint(self, typehint, registry):

sdks/python/apache_beam/dataframe/convert.py

@@ -17,10 +17,10 @@
 import inspect
 import warnings
 import weakref
-from typing import TYPE_CHECKING
 from typing import Any
 from typing import Dict
 from typing import Iterable
+from typing import Optional
 from typing import Tuple
 from typing import Union
 
@@ -35,19 +35,13 @@
 from apache_beam.dataframe.schemas import generate_proxy
 from apache_beam.typehints.pandas_type_compatibility import dtype_to_fieldtype
 
-if TYPE_CHECKING:
- # pylint: disable=ungrouped-imports
- from typing import Optional
-
 
 # TODO: Or should this be called as_dataframe?
 def to_dataframe(
- pcoll, # type: pvalue.PCollection
- proxy=None, # type: Optional[pd.core.generic.NDFrame]
- label=None, # type: Optional[str]
-):
- # type: (...) -> frame_base.DeferredFrame
-
+ pcoll: pvalue.PCollection,
+ proxy: Optional[pd.core.generic.NDFrame] = None,
+ label: Optional[str] = None,
+) -> frame_base.DeferredFrame:
  """Converts a PCollection to a deferred dataframe-like object, which can
  manipulated with pandas methods like `filter` and `groupby`.
 
@@ -93,10 +87,10 @@ def to_dataframe(
 # Note that the pipeline (indirectly) holds references to the transforms which
 # keeps both the PCollections and expressions alive. This ensures the
 # expression's ids are never accidentally re-used.
-TO_PCOLLECTION_CACHE = weakref.WeakValueDictionary(
-) # type: weakref.WeakValueDictionary[str, pvalue.PCollection]
-UNBATCHED_CACHE = weakref.WeakValueDictionary(
-) # type: weakref.WeakValueDictionary[str, pvalue.PCollection]
+TO_PCOLLECTION_CACHE: 'weakref.WeakValueDictionary[str, pvalue.PCollection]' = (
+  weakref.WeakValueDictionary())
+UNBATCHED_CACHE: 'weakref.WeakValueDictionary[str, pvalue.PCollection]' = (
+  weakref.WeakValueDictionary())
 
 
 class RowsToDataFrameFn(beam.DoFn):
@@ -173,7 +167,7 @@ def infer_output_type(self, input_element_type):
 
 
 def to_pcollection(
- *dataframes, # type: Union[frame_base.DeferredFrame, pd.DataFrame, pd.Series]
+ *dataframes: Union[frame_base.DeferredFrame, pd.DataFrame, pd.Series],
  label=None,
  always_return_tuple=False,
  yield_elements='schemas',
@@ -258,12 +252,12 @@ def extract_input(placeholder):
  df for df in dataframes if df._expr._id not in TO_PCOLLECTION_CACHE
  ]
  if len(new_dataframes):
- new_results = {p: extract_input(p)
-  for p in placeholders
-  } | label >> transforms._DataframeExpressionsTransform({
-  ix: df._expr
-  for (ix, df) in enumerate(new_dataframes)
-  }) # type: Dict[Any, pvalue.PCollection]
+ new_results: Dict[Any, pvalue.PCollection] = {
+ p: extract_input(p)
+ for p in placeholders
+ } | label >> transforms._DataframeExpressionsTransform(
+ {ix: df._expr
+ for (ix, df) in enumerate(new_dataframes)})
 
  TO_PCOLLECTION_CACHE.update(
  {new_dataframes[ix]._expr._id: pc

sdks/python/apache_beam/dataframe/frame_base.py

@@ -38,7 +38,7 @@
 
 class DeferredBase(object):
 
- _pandas_type_map = {} # type: Dict[Union[type, None], type]
+ _pandas_type_map: Dict[Union[type, None], type] = {}
 
  def __init__(self, expr):
  self._expr = expr
@@ -197,8 +197,8 @@ def _proxy_method(
  inplace=False,
  base=None,
  *,
- requires_partition_by, # type: partitionings.Partitioning
- preserves_partition_by, # type: partitionings.Partitioning
+ requires_partition_by: partitionings.Partitioning,
+ preserves_partition_by: partitionings.Partitioning,
 ):
  if name is None:
  name, func = name_and_func(func)
@@ -227,14 +227,14 @@ def _elementwise_function(
 
 
 def _proxy_function(
- func, # type: Union[Callable, str]
- name=None, # type: Optional[str]
- restrictions=None, # type: Optional[Dict[str, Union[Any, List[Any]]]]
- inplace=False, # type: bool
- base=None, # type: Optional[type]
+ func: Union[Callable, str],
+ name: Optional[str] = None,
+ restrictions: Optional[Dict[str, Union[Any, List[Any]]]] = None,
+ inplace: bool = False,
+ base: Optional[type] = None,
  *,
- requires_partition_by, # type: partitionings.Partitioning
- preserves_partition_by, # type: partitionings.Partitioning
+ requires_partition_by: partitionings.Partitioning,
+ preserves_partition_by: partitionings.Partitioning,
 ):
 
  if name is None:

sdks/python/apache_beam/dataframe/partitionings.py

@@ -32,9 +32,7 @@ class Partitioning(object):
  def __repr__(self):
  return self.__class__.__name__
 
- def is_subpartitioning_of(self, other):
- # type: (Partitioning) -> bool
-
+ def is_subpartitioning_of(self, other: 'Partitioning') -> bool:
  """Returns whether self is a sub-partition of other.
 
  Specifically, returns whether something partitioned by self is necissarily
@@ -48,9 +46,8 @@ def __lt__(self, other):
  def __le__(self, other):
  return not self.is_subpartitioning_of(other)
 
- def partition_fn(self, df, num_partitions):
- # type: (Frame, int) -> Iterable[Tuple[Any, Frame]]
-
+ def partition_fn(self, df: Frame,
+ num_partitions: int) -> Iterable[Tuple[Any, Frame]]:
  """A callable that actually performs the partitioning of a Frame df.
 
  This will be invoked via a FlatMap in conjunction with a GroupKey to

sdks/python/apache_beam/dataframe/schemas.py

@@ -85,9 +85,7 @@ def expand(self, pcoll):
  | beam.Map(converter.produce_batch))
 
 
-def generate_proxy(element_type):
- # type: (type) -> pd.DataFrame
-
+def generate_proxy(element_type: type) -> pd.DataFrame:
  """Generate a proxy pandas object for the given PCollection element_type.
 
  Currently only supports generating a DataFrame proxy from a schema-aware
@@ -106,9 +104,8 @@ def generate_proxy(element_type):
  return proxy
 
 
-def element_type_from_dataframe(proxy, include_indexes=False):
- # type: (pd.DataFrame, bool) -> type
-
+def element_type_from_dataframe(
+ proxy: pd.DataFrame, include_indexes: bool = False) -> type:
  """Generate an element_type for an element-wise PCollection from a proxy
  pandas object. Currently only supports converting the element_type for
  a schema-aware PCollection to a proxy DataFrame.