# Path: tensorflow_transform/beam/deep_copy.py
# def deep_copy(pcollection):
#   """Create a deep copy of a PCollection up to materialization boundaries."""
#   if not isinstance(pcollection, pvalue.PCollection):
#     raise ValueError('Input to deep_copy must be a PCollection.')
# 
#   # AppliedPTransform.update_input_refcounts() is a vestigial method that
#   # uses an incorrect heuristic; it will be removed in a future version of
#   # Beam, since its results aren't used anyway.  Until then, we work around
#   # this (see https://issues.apache.org/jira/browse/BEAM-4593).
#   if getattr(beam_pipeline.AppliedPTransform,
#              'update_input_refcounts', None) is not None:
#     beam_pipeline.AppliedPTransform.update_input_refcounts = lambda _: None
# 
#   to_clone = _get_items_to_clone(pcollection)
#   pcollection_replacements = _clone_items(pcollection.pipeline, to_clone)
# 
#   return pcollection_replacements[pcollection]
# 
# Path: tensorflow_transform/beam/test_helpers.py
# def make_test_beam_pipeline_kwargs():


import collections
import apache_beam as beam
import unittest
from apache_beam import pvalue
from apache_beam.transforms import resources
from tensorflow_transform.beam import deep_copy
from tensorflow_transform.beam import test_helpers
  @staticmethod
  def _MakeAdd1CountingIdentityFn(label):

    def Add1CountingIdentityFn(x_y):
      (x, y) = x_y
      return DeepCopyTest._CountingIdentityFn(label, (x + 1, y))

    return Add1CountingIdentityFn

  @staticmethod
  def _InitializeCounts():
    DeepCopyTest._counts = collections.defaultdict(int)

  def setUp(self):
    DeepCopyTest._InitializeCounts()

  def testBasicDeepCopy(self):
    with DeepCopyTest._MakeBeamPipeline() as p:
      grouped = (p
                 | beam.Create([(1, 'a'), (2, 'b'), (3, 'c')])
                 | beam.Map(
                     lambda x: DeepCopyTest._CountingIdentityFn(
                         'PreGroup', x))
                 | beam.GroupByKey())
      modified = (
          grouped
          |
          'Add1' >> beam.Map(DeepCopyTest._MakeAdd1CountingIdentityFn('Add1'))
          |
          'Add2' >> beam.Map(DeepCopyTest._MakeAdd1CountingIdentityFn('Add2')))

# Path: tensorflow_transform/beam/deep_copy.py
# def deep_copy(pcollection):
#   """Create a deep copy of a PCollection up to materialization boundaries."""
#   if not isinstance(pcollection, pvalue.PCollection):
#     raise ValueError('Input to deep_copy must be a PCollection.')
# 
#   # AppliedPTransform.update_input_refcounts() is a vestigial method that
#   # uses an incorrect heuristic; it will be removed in a future version of
#   # Beam, since its results aren't used anyway.  Until then, we work around
#   # this (see https://issues.apache.org/jira/browse/BEAM-4593).
#   if getattr(beam_pipeline.AppliedPTransform,
#              'update_input_refcounts', None) is not None:
#     beam_pipeline.AppliedPTransform.update_input_refcounts = lambda _: None
# 
#   to_clone = _get_items_to_clone(pcollection)
#   pcollection_replacements = _clone_items(pcollection.pipeline, to_clone)
# 
#   return pcollection_replacements[pcollection]
# 
# Path: tensorflow_transform/beam/test_helpers.py
# def make_test_beam_pipeline_kwargs():


# Path: tensorflow_transform/beam/deep_copy_test.py
import collections
import apache_beam as beam
import unittest
from apache_beam import pvalue
from apache_beam.transforms import resources
from tensorflow_transform.beam import deep_copy
from tensorflow_transform.beam import test_helpers

  @staticmethod
  def _MakeAdd1CountingIdentityFn(label):

    def Add1CountingIdentityFn(x_y):
      (x, y) = x_y
      return DeepCopyTest._CountingIdentityFn(label, (x + 1, y))

    return Add1CountingIdentityFn

  @staticmethod
  def _InitializeCounts():
    DeepCopyTest._counts = collections.defaultdict(int)

  def setUp(self):
    DeepCopyTest._InitializeCounts()

  def testBasicDeepCopy(self):
    with DeepCopyTest._MakeBeamPipeline() as p:
      grouped = (p
                 | beam.Create([(1, 'a'), (2, 'b'), (3, 'c')])
                 | beam.Map(
                     lambda x: DeepCopyTest._CountingIdentityFn(
                         'PreGroup', x))
                 | beam.GroupByKey())
      modified = (
          grouped
          |
          'Add1' >> beam.Map(DeepCopyTest._MakeAdd1CountingIdentityFn('Add1'))
          |
          'Add2' >> beam.Map(DeepCopyTest._MakeAdd1CountingIdentityFn('Add2')))

      copied = deep_copy.deep_copy(modified)