Merge pull request #15044 Coder and Windowing cleanups.

robertwb · web-flow · commit 8b66f37ccb3c · 2021-06-21T16:17:56.000-07:00
diff --git a/sdks/python/apache_beam/coders/coders.py b/sdks/python/apache_beam/coders/coders.py
@@ -220,7 +220,7 @@ def _dict_without_impl(self):
     return self.__dict__
 
   def to_type_hint(self):
-    raise NotImplementedError('BEAM-2717')
+    raise NotImplementedError('BEAM-2717: %s' % self.__class__.__name__)
 
   @classmethod
   def from_type_hint(cls, unused_typehint, unused_registry):
@@ -532,6 +532,9 @@ def __eq__(self, other):
   def __hash__(self):
     return hash(type(self)) + hash(self._key_coder) + hash(self._value_coder)
 
+  def __repr__(self):
+    return 'MapCoder[%s, %s]' % (self._key_coder, self._value_coder)
+
 
 class NullableCoder(FastCoder):
   def __init__(self, value_coder):
@@ -929,10 +932,10 @@ def __eq__(self, other):
   def __hash__(self):
     return hash(self.proto_message_type)
 
-  @staticmethod
-  def from_type_hint(typehint, unused_registry):
+  @classmethod
+  def from_type_hint(cls, typehint, unused_registry):
     if issubclass(typehint, proto_utils.message_types):
-      return ProtoCoder(typehint)
+      return cls(typehint)
     else:
       raise ValueError((
           'Expected a subclass of google.protobuf.message.Message'
@@ -1019,10 +1022,10 @@ def as_deterministic_coder(self, step_label, error_message=None):
   def to_type_hint(self):
     return typehints.Tuple[tuple(c.to_type_hint() for c in self._coders)]
 
-  @staticmethod
-  def from_type_hint(typehint, registry):
+  @classmethod
+  def from_type_hint(cls, typehint, registry):
     # type: (typehints.TupleConstraint, CoderRegistry) -> TupleCoder
-    return TupleCoder([registry.get_coder(t) for t in typehint.tuple_types])
+    return cls([registry.get_coder(t) for t in typehint.tuple_types])
 
   def as_cloud_object(self, coders_context=None):
     if self.is_kv_coder():
@@ -1106,10 +1109,10 @@ def as_deterministic_coder(self, step_label, error_message=None):
       return TupleSequenceCoder(
           self._elem_coder.as_deterministic_coder(step_label, error_message))
 
-  @staticmethod
-  def from_type_hint(typehint, registry):
+  @classmethod
+  def from_type_hint(cls, typehint, registry):
     # type: (Any, CoderRegistry) -> TupleSequenceCoder
-    return TupleSequenceCoder(registry.get_coder(typehint.inner_type))
+    return cls(registry.get_coder(typehint.inner_type))
 
   def _get_component_coders(self):
     # type: () -> Tuple[Coder, ...]
@@ -1484,11 +1487,11 @@ def to_type_hint(self):
     return sharded_key_type.ShardedKeyTypeConstraint(
         self._key_coder.to_type_hint())
 
-  @staticmethod
-  def from_type_hint(typehint, registry):
+  @classmethod
+  def from_type_hint(cls, typehint, registry):
     from apache_beam.typehints import sharded_key_type
     if isinstance(typehint, sharded_key_type.ShardedKeyTypeConstraint):
-      return ShardedKeyCoder(registry.get_coder(typehint.key_type))
+      return cls(registry.get_coder(typehint.key_type))
     else:
       raise ValueError((
           'Expected an instance of ShardedKeyTypeConstraint'
diff --git a/sdks/python/apache_beam/io/flink/flink_streaming_impulse_source.py b/sdks/python/apache_beam/io/flink/flink_streaming_impulse_source.py
@@ -42,7 +42,7 @@ def expand(self, pbegin):
         'Input to transform must be a PBegin but found %s' % pbegin)
     return pvalue.PCollection(pbegin.pipeline, is_bounded=False)
 
-  def get_windowing(self, inputs):
+  def get_windowing(self, unused_inputs):
     return Windowing(GlobalWindows())
 
   def infer_output_type(self, unused_input_type):
diff --git a/sdks/python/apache_beam/runners/direct/direct_runner.py b/sdks/python/apache_beam/runners/direct/direct_runner.py
@@ -398,7 +398,7 @@ def _infer_output_coder(
     # type: (...) -> typing.Optional[coders.Coder]
     return coders.BytesCoder()
 
-  def get_windowing(self, inputs):
+  def get_windowing(self, unused_inputs):
     return beam.Windowing(beam.window.GlobalWindows())
 
   def expand(self, pvalue):
diff --git a/sdks/python/apache_beam/transforms/core.py b/sdks/python/apache_beam/transforms/core.py
@@ -2887,13 +2887,6 @@ def expand(self, pcolls):
     is_bounded = all(pcoll.is_bounded for pcoll in pcolls)
     return pvalue.PCollection(self.pipeline, is_bounded=is_bounded)
 
-  def get_windowing(self, inputs):
-    # type: (typing.Any) -> Windowing
-    if not inputs:
-      # TODO(robertwb): Return something compatible with every windowing?
-      return Windowing(GlobalWindows())
-    return super(Flatten, self).get_windowing(inputs)
-
   def infer_output_type(self, input_type):
     return input_type
 
diff --git a/sdks/python/apache_beam/transforms/ptransform.py b/sdks/python/apache_beam/transforms/ptransform.py
@@ -549,8 +549,13 @@ def get_windowing(self, inputs):
     By default most transforms just return the windowing function associated
     with the input PCollection (or the first input if several).
     """
-    # TODO(robertwb): Assert all input WindowFns compatible.
-    return inputs[0].windowing
+    if inputs:
+      return inputs[0].windowing
+    else:
+      from apache_beam.transforms.core import Windowing
+      from apache_beam.transforms.window import GlobalWindows
+      # TODO(robertwb): Return something compatible with every windowing?
+      return Windowing(GlobalWindows())
 
   def __rrshift__(self, label):
     return _NamedPTransform(self, label)
