pyglove 0.4.5.dev202501140808__py3-none-any.whl → 0.4.5.dev202501150808__py3-none-any.whl

pyglove/core/utils/timing.py

@@ -56,6 +56,16 @@ class TimeIt:
           **kwargs,
       )
 
+    def merge(self, other: 'TimeIt.Status') -> 'TimeIt.Status':
+      """Merges the status of two `pg.timeit`."""
+      assert other.name == self.name, (self.name, other.name)
+      return TimeIt.Status(
+          name=self.name,
+          elapse=self.elapse + other.elapse,
+          has_ended=self.has_ended and other.has_ended,
+          error=self.error or other.error,
+      )
+
   @dataclasses.dataclass
   class StatusSummary(json_conversion.JSONConvertible):
     """Aggregated summary for repeated calls for `pg.timeit`."""
@@ -125,7 +135,7 @@ class TimeIt:
     self._name: str = name
     self._start_time: Optional[float] = None
     self._end_time: Optional[float] = None
-    self._child_contexts: Dict[str, TimeIt] = {}
+    self._child_contexts: List[TimeIt] = []
     self._error: Optional[error_utils.ErrorInfo] = None
     self._parent: Optional[TimeIt] = None
 
@@ -137,13 +147,11 @@ class TimeIt:
   @property
   def children(self) -> List['TimeIt']:
     """Returns child contexts."""
-    return list(self._child_contexts.values())
+    return self._child_contexts
 
   def add(self, context: 'TimeIt'):
     """Adds a child context."""
-    if context.name in self._child_contexts:
-      raise ValueError(f'`timeit` with name {context.name!r} already exists.')
-    self._child_contexts[context.name] = context
+    self._child_contexts.append(context)
 
   def start(self):
     """Starts timing."""
@@ -206,11 +214,14 @@ class TimeIt:
             has_ended=self.has_ended, error=self._error,
         )
     }
-    for child in self._child_contexts.values():
+    for child in self._child_contexts:
       child_result = child.status()
       for k, v in child_result.items():
         key = f'{self.name}.{k}' if self.name else k
-        result[key] = v
+        if key in result:
+          result[key] = result[key].merge(v)
+        else:
+          result[key] = v
     return result
 
   def __enter__(self):

pyglove/core/utils/timing_test.py

@@ -62,9 +62,9 @@ class TimeItTest(unittest.TestCase):
         self.assertFalse(r['node.child'].has_ended)
         self.assertTrue(r['node.child.grandchild'].has_ended)
 
-        with self.assertRaisesRegex(ValueError, '.* already exists'):
-          with timing.timeit('grandchild'):
-            pass
+        with timing.timeit('grandchild') as t3:
+          time.sleep(0.5)
+        self.assertEqual(t1.children, [t2, t3])
 
     elapse2 = t.elapse
     self.assertTrue(t.has_ended)
@@ -72,17 +72,20 @@ class TimeItTest(unittest.TestCase):
     time.sleep(0.5)
     self.assertEqual(elapse2, t.elapse)
 
-    statuss = t.status()
+    status = t.status()
     self.assertEqual(
-        list(statuss.keys()),
+        list(status.keys()),
         ['node', 'node.child', 'node.child.grandchild']
     )
     self.assertEqual(
-        sorted([v.elapse for v in statuss.values()], reverse=True),
-        [v.elapse for v in statuss.values()],
+        status['node.child.grandchild'].elapse, t2.elapse + t3.elapse
+    )
+    self.assertEqual(
+        sorted([v.elapse for v in status.values()], reverse=True),
+        [v.elapse for v in status.values()],
     )
-    self.assertTrue(all(v.has_ended for v in statuss.values()))
-    self.assertFalse(any(v.has_error for v in statuss.values()))
+    self.assertTrue(all(v.has_ended for v in status.values()))
+    self.assertFalse(any(v.has_error for v in status.values()))
     status = t.status()
     json_dict = json_conversion.to_json(status)
     status2 = json_conversion.from_json(json_dict)

{pyglove-0.4.5.dev202501140808.dist-info → pyglove-0.4.5.dev202501150808.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pyglove
-Version: 0.4.5.dev202501140808
+Version: 0.4.5.dev202501150808
 Summary: PyGlove: A library for manipulating Python objects.
 Home-page: https://github.com/google/pyglove
 Author: PyGlove Authors

{pyglove-0.4.5.dev202501140808.dist-info → pyglove-0.4.5.dev202501150808.dist-info}/RECORD

@@ -146,8 +146,8 @@ pyglove/core/utils/text_color.py,sha256=xcCTCxY2qFNZs_jismMGus8scEXKBpYGAhpAgnz-
 pyglove/core/utils/text_color_test.py,sha256=3HtZpUB5XPr7A_5Fg4ZSMfNWeDRiQgSzmg9b1tctMI4,2801
 pyglove/core/utils/thread_local.py,sha256=i-CnyY3VREtLfAj4_JndBnsKuQLIgwG29ma8dAyRxbI,4839
 pyglove/core/utils/thread_local_test.py,sha256=AOqsdNsA-cYMvJckqxb25ql3Y5kDW0qLfBW1cl85Bnw,6757
-pyglove/core/utils/timing.py,sha256=fXdwEYXR84qZBC7hxbXbwztW2UKliJZLqCDvhz1qw_A,7241
-pyglove/core/utils/timing_test.py,sha256=OZxi551MJec9cKtpNbeep6jjzP68JUHhSmRHNKDFRMw,4928
+pyglove/core/utils/timing.py,sha256=gnHCA2IJ9DZ3Y0vb8ItB1Ap8ss8j_zN0OMUAdhoGy2Y,7550
+pyglove/core/utils/timing_test.py,sha256=t_awsYs5SqEO3_2u6HDbVxRcSxSfHsE4zMbFG-1ErZw,5011
 pyglove/core/utils/value_location.py,sha256=wAryIwQeEfrRddyGRk-KbLA7dnNDLhL-dSyFI9wIG5U,26756
 pyglove/core/utils/value_location_test.py,sha256=X6Gih3IoYugziwXWH8VGz5bPeb3Kq0CfZxwbNVIsZJo,21338
 pyglove/core/views/__init__.py,sha256=gll9ZBRYz4p_-LWOdzSR2a6UTWcJ8nR430trrP0yLCU,967
@@ -207,8 +207,8 @@ pyglove/ext/scalars/randoms.py,sha256=LkMIIx7lOq_lvJvVS3BrgWGuWl7Pi91-lA-O8x_gZs
 pyglove/ext/scalars/randoms_test.py,sha256=nEhiqarg8l_5EOucp59CYrpO2uKxS1pe0hmBdZUzRNM,2000
 pyglove/ext/scalars/step_wise.py,sha256=IDw3tuTpv0KVh7AN44W43zqm1-E0HWPUlytWOQC9w3Y,3789
 pyglove/ext/scalars/step_wise_test.py,sha256=TL1vJ19xVx2t5HKuyIzGoogF7N3Rm8YhLE6JF7i0iy8,2540
-pyglove-0.4.5.dev202501140808.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-pyglove-0.4.5.dev202501140808.dist-info/METADATA,sha256=_HkFXhfIYMiYBTTw47VFpp0lnLW5yojjf9NomI1d3II,7067
-pyglove-0.4.5.dev202501140808.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-pyglove-0.4.5.dev202501140808.dist-info/top_level.txt,sha256=wITzJSKcj8GZUkbq-MvUQnFadkiuAv_qv5qQMw0fIow,8
-pyglove-0.4.5.dev202501140808.dist-info/RECORD,,
+pyglove-0.4.5.dev202501150808.dist-info/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+pyglove-0.4.5.dev202501150808.dist-info/METADATA,sha256=djfHeOlrh9nRfIY2ZQghYSlCVa_bBXxNOpjHCXPcpwE,7067
+pyglove-0.4.5.dev202501150808.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+pyglove-0.4.5.dev202501150808.dist-info/top_level.txt,sha256=wITzJSKcj8GZUkbq-MvUQnFadkiuAv_qv5qQMw0fIow,8
+pyglove-0.4.5.dev202501150808.dist-info/RECORD,,