From c970c730fc0457d8f20b02675c470bb57df2c56d Mon Sep 17 00:00:00 2001
From: JamesWrigley <james@puiterwijk.org>
Date: Sun, 16 Jul 2023 14:12:46 +0200
Subject: [PATCH] Implement .nbytes, .size_mb, and .size_gb for KeyData

Added because it can be useful to see how large a dataset will be in memory
without actually loading it.
---
 docs/reading_files.rst           |  6 ++++++
 extra_data/keydata.py            | 15 +++++++++++++++
 extra_data/tests/test_keydata.py |  3 +++
 3 files changed, 24 insertions(+)

diff --git a/docs/reading_files.rst b/docs/reading_files.rst
index 4ee5d8ba..1939c6f4 100644
--- a/docs/reading_files.rst
+++ b/docs/reading_files.rst
@@ -177,6 +177,12 @@ below, e.g.::
 
    .. automethod:: data_counts
 
+   .. autoattribute:: nbytes
+
+   .. autoattribute:: size_mb
+
+   .. autoattribute:: size_gb
+
    .. automethod:: ndarray
 
    .. automethod:: series
diff --git a/extra_data/keydata.py b/extra_data/keydata.py
index d2a00601..b121c113 100644
--- a/extra_data/keydata.py
+++ b/extra_data/keydata.py
@@ -91,6 +91,21 @@ def shape(self):
         """
         return (sum(c.total_count for c in self._data_chunks),) + self.entry_shape
 
+    @property
+    def nbytes(self):
+        """The number of bytes this data would take up in memory."""
+        return self.dtype.itemsize * np.prod(self.shape)
+
+    @property
+    def size_mb(self):
+        """The size of the data in memory in megabytes."""
+        return self.nbytes / 1e6
+
+    @property
+    def size_gb(self):
+        """The size of the data in memory in gigabytes."""
+        return self.nbytes / 1e9
+
     @property
     def source_file_paths(self):
         paths = []
diff --git a/extra_data/tests/test_keydata.py b/extra_data/tests/test_keydata.py
index 012be620..d2a9147a 100644
--- a/extra_data/tests/test_keydata.py
+++ b/extra_data/tests/test_keydata.py
@@ -34,6 +34,9 @@ def test_get_keydata(mock_spb_raw_run):
         'RAW-R0238-DA01-S00000.h5', 'RAW-R0238-DA01-S00001.h5'
     }
 
+    data = xgm_beam_x.ndarray()
+    assert xgm_beam_x.nbytes == data.nbytes
+
 def test_select_trains(mock_spb_raw_run):
     run = RunDirectory(mock_spb_raw_run)
     xgm_beam_x = run['SPB_XTD9_XGM/DOOCS/MAIN', 'beamPosition.ixPos.value']