# Path: tensorflow_transform/graph_context.py
# class TFGraphContext:
#   """A context manager to pass global state to a TF graph when it is traced.
# 
#   All the attributes in this context are kept on a thread local state.
# 
#   Attributes:
#     module_to_export: A tf.Module object that can be exported to a SavedModel
#       and will be used to track objects created within this TF graph.
#     temp_dir: The base path of the directory to write out any temporary files
#       in this context block. If None, the TF graph in this context will be
#       traced with placeholders for asset filepaths and is not serializable to a
#       SavedModel.
#     evaluated_replacements: A subset of placeholders/temporary asset files in
#       `analyzer_nodes.TENSOR_REPLACEMENTS` that have been evaluated in
#       previous TFT phases.
# 
#   Note that the temp dir should be accessible to worker jobs, e.g. if running
#   with the Cloud Dataflow runner, the temp dir should be on GCS and should have
#   permissions that allow both launcher and workers to access it.
#   """
# 
#   class _State(
#       tfx_namedtuple.namedtuple('_State', [
#           'module_to_export',
#           'temp_dir',
#           'evaluated_replacements',
#       ])):
#     """A named tuple storing state passed to this context manager."""
# 
#     @classmethod
#     def make_empty(cls):
#       """Return `_State` object with all fields set to `None`."""
#       return cls(*(None,) * len(cls._fields))
# 
#   _TEMP_SUBDIR = 'analyzer_temporary_assets'
# 
#   _thread_local = threading.local()
# 
#   def __init__(self,
#                module_to_export: tf.Module,
#                temp_dir: Optional[str] = None,
#                evaluated_replacements: Optional[Dict[str, Any]] = None):
#     self._module_to_export = module_to_export
#     self._temp_dir = temp_dir
#     self._evaluated_replacements = evaluated_replacements
# 
#   def __enter__(self):
#     assert getattr(self._thread_local, 'current_state', None) is None
#     self._thread_local.current_state = self._State(
#         module_to_export=self._module_to_export,
#         temp_dir=self._temp_dir,
#         evaluated_replacements=self._evaluated_replacements)
# 
#   def __exit__(self, *exn_info):
#     self._thread_local.current_state = None
# 
#   @property
#   def module_to_export(self):
#     return self._module_to_export
# 
#   @classmethod
#   def _get_current_state(cls) -> 'TFGraphContext._State':
#     if hasattr(cls._thread_local, 'current_state'):
#       return cls._thread_local.current_state
#     return cls._State.make_empty()
# 
#   @classmethod
#   def get_or_create_temp_dir(cls) -> Optional[str]:
#     """Generate a temporary location."""
#     current_state = cls._get_current_state()
#     if current_state.temp_dir is None:
#       return None
#     if not current_state.temp_dir:
#       raise ValueError('A temp dir was requested, but empty temp_dir was set. '
#                        'Use the TFGraphContext context manager.')
#     result = os.path.join(current_state.temp_dir, cls._TEMP_SUBDIR)
#     tf.io.gfile.makedirs(result)
#     return result
# 
#   @classmethod
#   def get_evaluated_replacements(cls) -> Optional[Dict[str, Any]]:
#     """Retrieves the value of evaluated_replacements if set.
# 
#     None otherwise.
# 
#     Returns:
#       A dictionary from graph tensor names to evaluated values for these
#       tensors. The keys are a subset of placeholders/temporary asset files in
#       `analyzer_nodes.TENSOR_REPLACEMENTS` that have been evaluated in
#       previous TFT phases.
#     """
#     return cls._get_current_state().evaluated_replacements
# 
#   @classmethod
#   def get_module_to_export(cls) -> Optional[tf.Module]:
#     """Retrieves the value of module_to_export.
# 
#     None if called outside a TFGraphContext scope.
# 
#     Returns:
#       A tf.Module object
#     """
#     return cls._get_current_state().module_to_export


import contextlib
import os
import tensorflow as tf
from typing import Callable, List, Optional
from tensorflow_transform.graph_context import TFGraphContext
from tensorflow.python.framework import func_graph
from tensorflow.python.framework import ops
from tensorflow.python.training.tracking import base
            'unique name was passed.')
      setattr(module, name, trackable_object)


# Thread-Hostile
@contextlib.contextmanager
def object_tracker_scope(object_tracker: ObjectTracker):
  """A context to manage trackable objects.

  Collects all trackable objects annotated using `track_object` within the body
  of its scope.

  Args:
    object_tracker: The passed in ObjectTracker object

  Yields:
    A scope in which the object_tracker is active.
  """
  global _OBJECT_TRACKER
  # Multiple nested object_tracker_scope calls are not expected.
  assert _OBJECT_TRACKER is None
  _OBJECT_TRACKER = object_tracker
  try:
    yield
  finally:
    _OBJECT_TRACKER = None


def _get_object(name: str) -> Optional[base.Trackable]:
  """If an object is being tracked using `name` return it, else None."""

# Path: tensorflow_transform/graph_context.py
# class TFGraphContext:
#   """A context manager to pass global state to a TF graph when it is traced.
# 
#   All the attributes in this context are kept on a thread local state.
# 
#   Attributes:
#     module_to_export: A tf.Module object that can be exported to a SavedModel
#       and will be used to track objects created within this TF graph.
#     temp_dir: The base path of the directory to write out any temporary files
#       in this context block. If None, the TF graph in this context will be
#       traced with placeholders for asset filepaths and is not serializable to a
#       SavedModel.
#     evaluated_replacements: A subset of placeholders/temporary asset files in
#       `analyzer_nodes.TENSOR_REPLACEMENTS` that have been evaluated in
#       previous TFT phases.
# 
#   Note that the temp dir should be accessible to worker jobs, e.g. if running
#   with the Cloud Dataflow runner, the temp dir should be on GCS and should have
#   permissions that allow both launcher and workers to access it.
#   """
# 
#   class _State(
#       tfx_namedtuple.namedtuple('_State', [
#           'module_to_export',
#           'temp_dir',
#           'evaluated_replacements',
#       ])):
#     """A named tuple storing state passed to this context manager."""
# 
#     @classmethod
#     def make_empty(cls):
#       """Return `_State` object with all fields set to `None`."""
#       return cls(*(None,) * len(cls._fields))
# 
#   _TEMP_SUBDIR = 'analyzer_temporary_assets'
# 
#   _thread_local = threading.local()
# 
#   def __init__(self,
#                module_to_export: tf.Module,
#                temp_dir: Optional[str] = None,
#                evaluated_replacements: Optional[Dict[str, Any]] = None):
#     self._module_to_export = module_to_export
#     self._temp_dir = temp_dir
#     self._evaluated_replacements = evaluated_replacements
# 
#   def __enter__(self):
#     assert getattr(self._thread_local, 'current_state', None) is None
#     self._thread_local.current_state = self._State(
#         module_to_export=self._module_to_export,
#         temp_dir=self._temp_dir,
#         evaluated_replacements=self._evaluated_replacements)
# 
#   def __exit__(self, *exn_info):
#     self._thread_local.current_state = None
# 
#   @property
#   def module_to_export(self):
#     return self._module_to_export
# 
#   @classmethod
#   def _get_current_state(cls) -> 'TFGraphContext._State':
#     if hasattr(cls._thread_local, 'current_state'):
#       return cls._thread_local.current_state
#     return cls._State.make_empty()
# 
#   @classmethod
#   def get_or_create_temp_dir(cls) -> Optional[str]:
#     """Generate a temporary location."""
#     current_state = cls._get_current_state()
#     if current_state.temp_dir is None:
#       return None
#     if not current_state.temp_dir:
#       raise ValueError('A temp dir was requested, but empty temp_dir was set. '
#                        'Use the TFGraphContext context manager.')
#     result = os.path.join(current_state.temp_dir, cls._TEMP_SUBDIR)
#     tf.io.gfile.makedirs(result)
#     return result
# 
#   @classmethod
#   def get_evaluated_replacements(cls) -> Optional[Dict[str, Any]]:
#     """Retrieves the value of evaluated_replacements if set.
# 
#     None otherwise.
# 
#     Returns:
#       A dictionary from graph tensor names to evaluated values for these
#       tensors. The keys are a subset of placeholders/temporary asset files in
#       `analyzer_nodes.TENSOR_REPLACEMENTS` that have been evaluated in
#       previous TFT phases.
#     """
#     return cls._get_current_state().evaluated_replacements
# 
#   @classmethod
#   def get_module_to_export(cls) -> Optional[tf.Module]:
#     """Retrieves the value of module_to_export.
# 
#     None if called outside a TFGraphContext scope.
# 
#     Returns:
#       A tf.Module object
#     """
#     return cls._get_current_state().module_to_export


# Path: tensorflow_transform/annotators.py
import contextlib
import os
import tensorflow as tf
from typing import Callable, List, Optional
from tensorflow_transform.graph_context import TFGraphContext
from tensorflow.python.framework import func_graph
from tensorflow.python.framework import ops
from tensorflow.python.training.tracking import base

            'unique name was passed.')
      setattr(module, name, trackable_object)


# Thread-Hostile
@contextlib.contextmanager
def object_tracker_scope(object_tracker: ObjectTracker):
  """A context to manage trackable objects.

  Collects all trackable objects annotated using `track_object` within the body
  of its scope.

  Args:
    object_tracker: The passed in ObjectTracker object

  Yields:
    A scope in which the object_tracker is active.
  """
  global _OBJECT_TRACKER
  # Multiple nested object_tracker_scope calls are not expected.
  assert _OBJECT_TRACKER is None
  _OBJECT_TRACKER = object_tracker
  try:
    yield
  finally:
    _OBJECT_TRACKER = None


def _get_object(name: str) -> Optional[base.Trackable]:
  """If an object is being tracked using `name` return it, else None."""

  module = TFGraphContext.get_module_to_export()