try:
long # Python 2
except NameError:
long = int # Python 3
__all__ = [
'Count',
'Mean',
'Sample',
'Top',
'ToDict',
'ToList',
]
# Type variables
T = TypeVariable('T')
K = TypeVariable('K')
V = TypeVariable('V')
class Mean(object):
"""Combiners for computing arithmetic means of elements."""
class Globally(ptransform.PTransform):
"""combiners.Mean.Globally computes the arithmetic mean of the elements."""
def expand(self, pcoll):
return pcoll | core.CombineGlobally(MeanCombineFn())
class PerKey(ptransform.PTransform):
"""combiners.Mean.PerKey finds the means of the values for each key."""
def expand(self, pcoll):
return pcoll | core.CombinePerKey(MeanCombineFn())
import typing
import unittest
# patches unittest.TestCase to be python3 compatible
import future.tests.base # pylint: disable=unused-import
from apache_beam.typehints import Any
from apache_beam.typehints import Dict
from apache_beam.typehints import List
from apache_beam.typehints import Tuple
from apache_beam.typehints import TypeVariable
from apache_beam.typehints import decorators
decorators._enable_from_callable = True
T = TypeVariable('T')
# Name is 'T' so it converts to a beam type with the same name.
T_typing = typing.TypeVar('T')
class IOTypeHintsTest(unittest.TestCase):
def test_from_callable(self):
def fn(a: int, b: str = None, *args: Tuple[T], foo: List[int],
**kwargs: Dict[str, str]) -> Tuple[Any, ...]:
return a, b, args, foo, kwargs
th = decorators.IOTypeHints.from_callable(fn)
self.assertEqual(th.input_types, (
(int, str, Tuple[T]), {'foo': List[int], 'kwargs': Dict[str, str]}))
self.assertEqual(th.output_types, ((Tuple[Any, ...],), {}))