content: import contextlib
  import inspect
  import logging
  import os
  import os.path
  import shutil
  import tempfile
  import time
  import traceback
  import typing as t
  from contextvars import ContextVar
  from dataclasses import dataclass
  
  TWrite = t.Union[str, bytes]
  WriteObj = t.Union[TWrite, t.Callable[[], TWrite]]
  
  _CURRENT_WRITER: ContextVar["DiagnosticsWriter"] = ContextVar("_CURRENT_WRITER")
  _CURRENT_COLLECTOR: ContextVar["DiagnosticsCollector"] = ContextVar(
      "_CURRENT_COLLECTOR"
  )
  # Allows a collector to indicate subsequent collections should be suppressed to
  # avoid duplicate collections.
  _SUBSEQUENT_COLLECT_SUPPRESSED_BY: ContextVar[object] = ContextVar(
      "_SUBSEQUENT_COLLECT_SUPPRESSED_BY"
  )
  # Indicates current execution context is within a context manager by
  # `collect_when`. Only when it's set do we actually write diagnostics.
  _IS_IN_COLLECT_CONTEXT: ContextVar[bool] = ContextVar("_IS_IN_COLLECT_CONTEXT")
  _LOGGER = logging.getLogger(__name__)
  
  
  @dataclass
  class CollectionConditionContext:
      exception: t.Optional[Exception]
  
  
  CollectionCondition = t.Callable[[CollectionConditionContext], bool]
  
  
  def collect_when(
      condition: "CollectionCondition", supress_subsequent_collect: bool = True
  ):
      """See `DiagnosticsCollector.collect_when`"""
      return get_current_collector().collect_when(condition, supress_subsequent_collect)
  
  
  def collect():
      return collect_when(CollectionConditions.always())
  
  
  def collect_when_fail():
      return collect_when(CollectionConditions.when_fail())
  
  
  def write(file_name: str, text: WriteObj):
      return get_current_writer().write(file_name, text)
  
  
  def get_current_writer() -> "DiagnosticsWriter":
      """Get the writer for current execution context.
  
      Lazily instantiates and registers one if not already done.
      """
      current_writer = _CURRENT_WRITER.get(None)
      if not current_writer:
          current_writer = DiagnosticsWriter()
          _CURRENT_WRITER.set(current_writer)
      return current_writer
  
  
  def get_current_collector() -> "DiagnosticsCollector":
      current_collector = _CURRENT_COLLECTOR.get(None)
      if not current_collector:
          current_collector = DiagnosticsCollector()
          _CURRENT_COLLECTOR.set(current_collector)
      return current_collector
  
  
  def set_current_collector(collector: "DiagnosticsCollector"):
      _CURRENT_COLLECTOR.set(collector)
  
  
  class DiagnosticsWriter:
      # the root dir in which the diagnostics will be written
      _root_dir: str
  
      def __init__(self):
          self._root_dir = tempfile.mkdtemp(prefix="fx2trt.")
          self._data = ""
          _LOGGER.info(f"Initializing DiagnosticsWriter with root_dir: {self._root_dir}")
  
      def write(self, file_name: str, data: WriteObj):
          """
          TODO: Can be disabled by regex on file_name
          """
          self._data = data
          # Only write if we are inside a collect_when() context.
          if not _IS_IN_COLLECT_CONTEXT.get(False):
              return
  
          try:
              res, err = _res_or_err(data)
              if err:
                  to_write = err.encode("utf-8")
              else:
                  if isinstance(res, str):
                      to_write = res.encode("utf-8")
                  elif isinstance(res, bytes):
                      to_write = res
                  else:
                      raise TypeError(f"Unknown data type: {type(res)}")
              self._write(file_name, to_write)
          except Exception as e:
              # Log the error and swallow the exception, as this should not
              # propagated into business logic
              _LOGGER.warning(f"Error writing diagnostics: {e}")
  
      def root_dir(self) -> str:
          return self._root_dir
  
      def data(self) -> WriteObj:
          return self._data
  
      def _write(self, file_name: str, to_write: bytes):
          # ms granularity - no naming collash, otherwise file will be
          # overwritten.
          ts = int(time.time() * 1000)
          file_name = f"{file_name}.{ts}"
          fn = os.path.join(self.root_dir(), file_name)
          with open(fn, "wb") as f:
              f.write(to_write)
  
  
  class CollectionConditions:
      @classmethod
      def any(cls, *conditions: "CollectionCondition") -> "CollectionCondition":
          return lambda ctx: any(cond(ctx) for cond in conditions)
  
      @classmethod
      def all(cls, *conditions: "CollectionCondition") -> "CollectionCondition":
          return lambda ctx: all(cond(ctx) for cond in conditions)
  
      @classmethod
      def not_(cls, condition: "CollectionCondition") -> "CollectionCondition":
          return lambda ctx: not condition(ctx)
  
      @classmethod
      def always(cls) -> "CollectionCondition":
          """Always collect"""
          return lambda ctx: True
  
      @classmethod
      def never(cls) -> "CollectionCondition":
          """Never collect"""
          return lambda ctx: False
  
      @classmethod
      def when_fail(cls) -> "CollectionCondition":
          """Collect when failed"""
          ctx: CollectionConditionContext
          return lambda ctx: ctx.exception is not None
  
      @classmethod
      def when_called_by_function(
          cls, func_name: str, match_prefix: bool = False
      ) -> "CollectionCondition":
          def _when_called_by_function(ctx: CollectionConditionContext) -> bool:
              frames = inspect.stack()
              for frame in frames:
                  if match_prefix:
                      if frame[3].startswith(func_name):
                          return True
                  else:
                      if frame[3] == func_name:
                          return True
              return False
  
          return _when_called_by_function
  
      @classmethod
      def when_not_in_tests(cls) -> CollectionCondition:
          return CollectionConditions.not_(
              CollectionConditions.when_called_by_function("test_", match_prefix=True)
          )
  
  
  class DiagnosticsCollector:
      @contextlib.contextmanager
      def collect_when(
          self, condition: "CollectionCondition", supress_subsequent_collect: bool = True
      ):
          """
          Context manager to collect diagnostics when the enclosed code completes
          and *any* of the given condition is met.
  
          Args:
              condition:
                  the condition only when met should the collection be done
              supress_subsequent_collect:
                  When true, suppress any collections registered by this function
                  call. This is to ensure duplicate collections registered across
                  the callstack by different components. In this case, only the
                  outermost component will collect.
  
                  When false, always collect (subject to given condition) regardless
                  of earlier collection registration's suppression.
  
          Returns:
              a context manager that handles the collection when its enclosed
              code finished run.
          """
          this_collection_handle = object()
          suppressed_by = _SUBSEQUENT_COLLECT_SUPPRESSED_BY.get(None)
          reset_suppressed_by = False
          if supress_subsequent_collect:
              if suppressed_by and suppressed_by != this_collection_handle:
                  # Disable this collection since it's suppressed by a previously
                  # installed collection
                  condition = CollectionConditions.never()
              else:
                  suppressed_by = this_collection_handle
                  _SUBSEQUENT_COLLECT_SUPPRESSED_BY.set(suppressed_by)
                  # don't forget to reset it in `finanlly`
                  reset_suppressed_by = True
  
          is_in_collect_context_tok = _IS_IN_COLLECT_CONTEXT.set(True)
          exception: t.Optional[Exception] = None
          try:
              yield
          except Exception as e:
              exception = e
              raise
          finally:
              if reset_suppressed_by:
                  _SUBSEQUENT_COLLECT_SUPPRESSED_BY.set(None)
              if self._test_condition(condition, CollectionConditionContext(exception)):
                  try:
                      self.collect()
                  except Exception as e:
                      _LOGGER.warning(
                          f"Error while collecting diagnostics (THIS EXCEPTION IS HANDLED):\n"
                          f"{e}\n"
                          f"{traceback.format_exc()}"
                      )
              _IS_IN_COLLECT_CONTEXT.reset(is_in_collect_context_tok)
  
      def collect(self) -> str:
          """Collect the diagnostics. Overridable in sub-classes."""
          return ""
  
      @classmethod
      def _test_condition(
          cls, cond: CollectionCondition, ctx: CollectionConditionContext
      ) -> bool:
          try:
              return cond(ctx)
          except Exception as e:
              _LOGGER.warning(f"Error while testing condition: {e}")
              return False
  
  
  class ZipDiagnosticsCollector(DiagnosticsCollector):
      _write: DiagnosticsWriter
      _last_zip_path_for_test: str = ""  # for test purpose only
  
      def __init__(self, writer: DiagnosticsWriter):
          self._write = writer
  
      def collect(self) -> str:
          _, fp = tempfile.mkstemp()
          try:
              zip_path = shutil.make_archive(fp, "zip", self._write.root_dir())
              self._last_zip_path_for_test = zip_path
              return zip_path
          finally:
              os.remove(fp)
  
      def data(self) -> WriteObj:
          return self._write.data()
  
  
  def _res_or_err(data: WriteObj) -> t.Tuple[TWrite, str]:
      if isinstance(data, (str, bytes)):
          return data, ""
      if not callable(data):
          raise TypeError(
              f"data must be a callable that returns actual data to"
              f"write, but got {type(data)}"
          )
      try:
          return data(), ""
      except Exception as e:
          _LOGGER.warning(f"Error getting data to write: {e}")
          return "", str(e)
  