No public description

PiperOrigin-RevId: 671905636
google · Sep 7, 2024 · 9c34b92 · 9c34b92
1 parent 0d0cdf9
commit 9c34b92
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 31 deletions.
diff --git a/xarray_beam/_src/dataset.py b/xarray_beam/_src/dataset.py
@@ -115,7 +115,8 @@ def pipe(self, func, *args, **kwargs):
   def __repr__(self):
     base = repr(self.template)
     chunks_str = ', '.join(f'{k}: {v}' for k, v in self.chunks.items())
-    return base.replace(
-        '<xarray.Dataset>',
-        f'<xarray_beam.Dataset[{chunks_str}][split_vars={self.split_vars}]>',
+    return (
+        f'<xarray_beam.Dataset[{chunks_str}][split_vars={self.split_vars}]>'
+        + '\n'
+        + '\n'.join(base.split('\n')[1:])
     )
diff --git a/xarray_beam/_src/dataset_test.py b/xarray_beam/_src/dataset_test.py
@@ -34,14 +34,8 @@ def test_from_xarray(self):
     self.assertFalse(beam_ds.split_vars)
     self.assertRegex(beam_ds.ptransform.label, r'^from_xarray_\d+$')
     self.assertEqual(
-        repr(beam_ds),
-        textwrap.dedent("""
-            <xarray_beam.Dataset[x: 5][split_vars=False]>
-            Dimensions:  (x: 10)
-            Dimensions without coordinates: x
-            Data variables:
-                foo      (x) int64 dask.array<chunksize=(10,), meta=np.ndarray>
-        """).strip(),
+        repr(beam_ds).split('\n')[0],
+        "<xarray_beam.Dataset[x: 5][split_vars=False]>",
     )
     expected = [
         (xbeam.Key({'x': 0}), ds.head(x=5)),

diff --git a/xarray_beam/_src/rechunk_test.py b/xarray_beam/_src/rechunk_test.py
@@ -367,17 +367,6 @@ def test_consolidate_with_unchunked_vars(self):
                 combining nested dataset chunks for vars=None with offsets={'x': [0, 10]} failed.
                 Leading datasets along dimension 'x':
                   <xarray.Dataset>
-                  Dimensions:  (x: 10)
-                  Dimensions without coordinates: x
-                  Data variables:
-                      foo      (x) int64 0 1 2 3 4 5 6 7 8 9
-                      bar      int64 1
-                  <xarray.Dataset>
-                  Dimensions:  (x: 10)
-                  Dimensions without coordinates: x
-                  Data variables:
-                      foo      (x) int64 10 11 12 13 14 15 16 17 18 19
-                      bar      int64 2
                 """
             ).strip()
         ),
@@ -465,15 +454,6 @@ def test_consolidate_variables_merge_fails(self):
                 """
             merging dataset chunks with variables [{'foo'}, {'bar'}] failed.
               <xarray.Dataset>
-              Dimensions:  (x: 2)
-              Dimensions without coordinates: x
-              Data variables:
-                  foo      (x) int64 1 2
-              <xarray.Dataset>
-              Dimensions:  (x: 3)
-              Dimensions without coordinates: x
-              Data variables:
-                  bar      (x) int64 3 4 5
         """
             ).strip()
         ),