cleanup: replace namedtuple with dataclass - part 1 (#5998)

See #5725 for context.

Namedtuples used in `manager.py` is removed in a #6003.

Author: yatbear

tensorboard/backend/event_processing/event_accumulator.py

@@ -15,8 +15,11 @@
 """Takes a generator of values, and accumulates them for a frontend."""
 
 import collections
+import dataclasses
 import threading
 
+from typing import Sequence, Tuple
+
 from tensorboard.backend.event_processing import directory_watcher
 from tensorboard.backend.event_processing import event_file_loader
 from tensorboard.backend.event_processing import io_wrapper
@@ -27,47 +30,137 @@
 from tensorboard.compat.proto import event_pb2
 from tensorboard.compat.proto import graph_pb2
 from tensorboard.compat.proto import meta_graph_pb2
+from tensorboard.compat.proto import tensor_pb2
 from tensorboard.plugins.distribution import compressor
 from tensorboard.util import tb_logging
 
 
 logger = tb_logging.get_logger()
 
-namedtuple = collections.namedtuple
-ScalarEvent = namedtuple("ScalarEvent", ["wall_time", "step", "value"])
-
-CompressedHistogramEvent = namedtuple(
-    "CompressedHistogramEvent",
-    ["wall_time", "step", "compressed_histogram_values"],
-)
-
-HistogramEvent = namedtuple(
-    "HistogramEvent", ["wall_time", "step", "histogram_value"]
-)
-
-HistogramValue = namedtuple(
-    "HistogramValue",
-    ["min", "max", "num", "sum", "sum_squares", "bucket_limit", "bucket"],
-)
-
-ImageEvent = namedtuple(
-    "ImageEvent",
-    ["wall_time", "step", "encoded_image_string", "width", "height"],
-)
-
-AudioEvent = namedtuple(
-    "AudioEvent",
-    [
-        "wall_time",
-        "step",
-        "encoded_audio_string",
-        "content_type",
-        "sample_rate",
-        "length_frames",
-    ],
-)
-
-TensorEvent = namedtuple("TensorEvent", ["wall_time", "step", "tensor_proto"])
+
+@dataclasses.dataclass(frozen=True)
+class ScalarEvent:
+    """Contains information of a scalar event.
+
+    Attributes:
+      wall_time: Timestamp of the event in seconds.
+      step: Global step of the event.
+      value: A float or int value of the scalar.
+    """
+
+    wall_time: float
+    step: int
+    value: float
+
+
+@dataclasses.dataclass(frozen=True)
+class CompressedHistogramEvent:
+    """Contains information of a compressed histogram event.
+
+    Attributes:
+      wall_time: Timestamp of the event in seconds.
+      step: Global step of the event.
+      compressed_histogram_values: A sequence of tuples of basis points and
+        associated values in a compressed histogram.
+    """
+
+    wall_time: float
+    step: int
+    compressed_histogram_values: Sequence[Tuple[float, float]]
+
+
+@dataclasses.dataclass(frozen=True)
+class HistogramValue:
+    """Holds the information of the histogram values.
+
+    Attributes:
+      min: A float or int min value.
+      max: A float or int max value.
+      num: Total number of values.
+      sum: Sum of all values.
+      sum_squares: Sum of squares for all values.
+      bucket_limit: Upper values per bucket.
+      bucket: Numbers of values per bucket.
+    """
+
+    min: float
+    max: float
+    num: int
+    sum: float
+    sum_squares: float
+    bucket_limit: Sequence[float]
+    bucket: Sequence[int]
+
+
+@dataclasses.dataclass(frozen=True)
+class HistogramEvent:
+    """Contains information of a histogram event.
+
+    Attributes:
+      wall_time: Timestamp of the event in seconds.
+      step: Global step of the event.
+      histogram_value: Information of the histogram values.
+    """
+
+    wall_time: float
+    step: int
+    histogram_value: HistogramValue
+
+
+@dataclasses.dataclass(frozen=True)
+class ImageEvent:
+    """Contains information of an image event.
+
+    Attributes:
+      wall_time: Timestamp of the event in seconds.
+      step: Global step of the event.
+      encoded_image_string: Image content encoded in bytes.
+      width: Width of the image.
+      height: Height of the image.
+    """
+
+    wall_time: float
+    step: int
+    encoded_image_string: bytes
+    width: int
+    height: int
+
+
+@dataclasses.dataclass(frozen=True)
+class AudioEvent:
+    """Contains information of an audio event.
+
+    Attributes:
+      wall_time: Timestamp of the event in seconds.
+      step: Global step of the event.
+      encoded_audio_string: Audio content encoded in bytes.
+      content_type: A string describes the type of the audio content.
+      sample_rate: Sample rate of the audio in Hz. Must be positive.
+      length_frames: Length of the audio in frames (samples per channel).
+    """
+
+    wall_time: float
+    step: int
+    encoded_audio_string: bytes
+    content_type: str
+    sample_rate: float
+    length_frames: int
+
+
+@dataclasses.dataclass(frozen=True)
+class TensorEvent:
+    """A tensor event.
+
+    Attributes:
+      wall_time: Timestamp of the event in seconds.
+      step: Global step of the event.
+      tensor_proto: A `TensorProto`.
+    """
+
+    wall_time: float
+    step: int
+    tensor_proto: tensor_pb2.TensorProto
+
 
 ## Different types of summary events handled by the event_accumulator
 SUMMARY_TYPES = {
@@ -664,7 +757,8 @@ def _CheckForOutOfOrderStepAndMaybePurge(self, event):
             self.most_recent_step = event.step
             self.most_recent_wall_time = event.wall_time
 
-    def _ConvertHistogramProtoToTuple(self, histo):
+    def _ConvertHistogramProtoToPopo(self, histo):
+        """Converts histogram proto to Python object."""
         return HistogramValue(
             min=histo.min,
             max=histo.max,
@@ -677,7 +771,7 @@ def _ConvertHistogramProtoToTuple(self, histo):
 
     def _ProcessHistogram(self, tag, wall_time, step, histo):
         """Processes a proto histogram by adding it to accumulated state."""
-        histo = self._ConvertHistogramProtoToTuple(histo)
+        histo = self._ConvertHistogramProtoToPopo(histo)
         histo_ev = HistogramEvent(wall_time, step, histo)
         self.histograms.AddItem(tag, histo_ev)
         self.compressed_histograms.AddItem(

tensorboard/backend/event_processing/event_accumulator_test.py

@@ -374,7 +374,7 @@ def testCompressedHistograms(self):
 
         # Create the expected values after compressing hst1
         expected_vals1 = [
-            compressor.CompressedHistogramValue(bp, val)
+            compressor.CompressedHistogramValue(bp, val).as_tuple()
             for bp, val in [
                 (0, 1.0),
                 (2500, 1.25),
@@ -390,7 +390,7 @@ def testCompressedHistograms(self):
 
         # Create the expected values after compressing hst2
         expected_vals2 = [
-            compressor.CompressedHistogramValue(bp, val)
+            compressor.CompressedHistogramValue(bp, val).as_tuple()
             for bp, val in [
                 (0, -2),
                 (2500, 2),

tensorboard/backend/event_processing/event_file_inspector.py

@@ -109,10 +109,12 @@
 """
 
 
-import collections
+import dataclasses
 import itertools
 import os
 
+from typing import Any, Generator, Mapping
+
 from tensorboard.backend.event_processing import event_accumulator
 from tensorboard.backend.event_processing import event_file_loader
 from tensorboard.backend.event_processing import io_wrapper
@@ -145,25 +147,44 @@
 # All summary types that we can inspect.
 TRACKED_FIELDS = SHORT_FIELDS + LONG_FIELDS
 
-# An `Observation` contains the data within each Event file that the inspector
-# cares about. The inspector accumulates Observations as it processes events.
-Observation = collections.namedtuple(
-    "Observation", ["step", "wall_time", "tag"]
-)
+PRINT_SEPARATOR = "=" * 70 + "\n"
+
 
-# An InspectionUnit is created for each organizational structure in the event
-# files visible in the final terminal output. For instance, one InspectionUnit
-# is created for each subdirectory in logdir. When asked to inspect a single
-# event file, there may only be one InspectionUnit.
+@dataclasses.dataclass(frozen=True)
+class Observation:
+    """Contains the data within each Event file that the inspector cares about.
 
-# The InspectionUnit contains the `name` of the organizational unit that will be
-# printed to console, a `generator` that yields `Event` protos, and a mapping
-# from string fields to `Observations` that the inspector creates.
-InspectionUnit = collections.namedtuple(
-    "InspectionUnit", ["name", "generator", "field_to_obs"]
-)
+    The inspector accumulates Observations as it processes events.
+
+    Attributes:
+      step: Global step of the event.
+      wall_time: Timestamp of the event in seconds.
+      tag: Tag name associated with the event.
+    """
+
+    step: int
+    wall_time: float
+    tag: str
 
-PRINT_SEPARATOR = "=" * 70 + "\n"
+
+@dataclasses.dataclass(frozen=True)
+class InspectionUnit:
+    """Created for each organizational structure in the event files.
+
+    An InspectionUnit is visible in the final terminal output. For instance, one
+    InspectionUnit is created for each subdirectory in logdir. When asked to inspect
+    a single event file, there may only be one InspectionUnit.
+
+    Attributes:
+      name: Name of the organizational unit that will be printed to console.
+      generator: A generator that yields `Event` protos.
+      field_to_obs: A mapping from string fields to `Observations` that the inspector
+        creates.
+    """
+
+    name: str
+    generator: Generator[event_pb2.Event, Any, Any]
+    field_to_obs: Mapping[str, Observation]
 
 
 def get_field_to_observations_map(generator, query_for_tag=""):
@@ -181,9 +202,9 @@ def get_field_to_observations_map(generator, query_for_tag=""):
     def increment(stat, event, tag=""):
         assert stat in TRACKED_FIELDS
         field_to_obs[stat].append(
-            Observation(
-                step=event.step, wall_time=event.wall_time, tag=tag
-            )._asdict()
+            dataclasses.asdict(
+                Observation(step=event.step, wall_time=event.wall_time, tag=tag)
+            )
         )
 
     field_to_obs = dict([(t, []) for t in TRACKED_FIELDS])

tensorboard/backend/event_processing/plugin_event_accumulator.py

@@ -15,9 +15,9 @@
 """Takes a generator of values, and accumulates them for a frontend."""
 
 import collections
+import dataclasses
 import threading
 
-
 from tensorboard.backend.event_processing import directory_loader
 from tensorboard.backend.event_processing import directory_watcher
 from tensorboard.backend.event_processing import event_file_loader
@@ -29,14 +29,12 @@
 from tensorboard.compat.proto import event_pb2
 from tensorboard.compat.proto import graph_pb2
 from tensorboard.compat.proto import meta_graph_pb2
+from tensorboard.compat.proto import tensor_pb2
 from tensorboard.util import tb_logging
 
 
 logger = tb_logging.get_logger()
 
-namedtuple = collections.namedtuple
-
-TensorEvent = namedtuple("TensorEvent", ["wall_time", "step", "tensor_proto"])
 
 # Legacy aliases
 TENSORS = tag_types.TENSORS
@@ -55,6 +53,21 @@
 _TENSOR_RESERVOIR_KEY = "."  # arbitrary
 
 
+@dataclasses.dataclass(frozen=True)
+class TensorEvent:
+    """A tensor event.
+
+    Attributes:
+      wall_time: Timestamp of the event in seconds.
+      step: Global step of the event.
+      tensor_proto: A `TensorProto`.
+    """
+
+    wall_time: float
+    step: int
+    tensor_proto: tensor_pb2.TensorProto
+
+
 class EventAccumulator(object):
     """An `EventAccumulator` takes an event generator, and accumulates the
     values.