From 438980120b6e1075f1664c250648b61ddd5bf1c4 Mon Sep 17 00:00:00 2001
From: ctena <carles.tena@bsc.es>
Date: Sat, 27 Apr 2024 19:12:58 +0200
Subject: [PATCH] Direct access to variable data

---
 CHANGELOG.rst                     |  1 +
 nes/nc_projections/default_nes.py | 26 +++++++++++++++++++++++++-
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index b163242..f94241b 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -17,6 +17,7 @@ CHANGELOG
   * Documentation
   * Removed negative values on the horizontal interpolation due to unmapped NaNs values.
   * Improved load_nes.py removing redundant code
+  * Direct access to variable data. (`#77 <https://earth.bsc.es/gitlab/es/NES/-/issues/77>`_)
   * New functionalities for vertical extrapolation. (`#74 <https://earth.bsc.es/gitlab/es/NES/-/issues/74>`_)
   * Bugfix:
 
diff --git a/nes/nc_projections/default_nes.py b/nes/nc_projections/default_nes.py
index 4d07272..e3fb5ec 100644
--- a/nes/nc_projections/default_nes.py
+++ b/nes/nc_projections/default_nes.py
@@ -15,6 +15,7 @@ from shapely.geometry import Polygon, Point
 from copy import deepcopy, copy
 import datetime
 from dateutil.relativedelta import relativedelta
+from typing import Union, List
 import pyproj
 from ..methods import vertical_interpolation, horizontal_interpolation, cell_measures, spatial_join
 from ..nes_formats import to_netcdf_cams_ra, to_netcdf_monarch, to_monarch_units, to_netcdf_cmaq, to_cmaq_units, \
@@ -446,7 +447,30 @@ class Nes(object):
         else:
             return self.__add__(other)
 
-    def copy(self, copy_vars=False):
+    def __getitem__(self, key: str) -> Union[np.array, None]:
+        """
+        Retrieve the data associated with the specified key.
+
+        Parameters
+        ----------
+        key : str
+            The key to retrieve the data for.
+
+        Returns
+        -------
+        Union[np.array, None]
+            The data associated with the specified key, or None if the key
+            does not exist.
+
+        Notes
+        -----
+        This method allows accessing data in the variables dictionary using
+        dictionary-like syntax, e.g., obj[key]['data'].
+
+        """
+        return self.variables[key]['data']
+
+    def copy(self, copy_vars: bool = False):
         """
         Copy the Nes object.
         The copy will avoid to copy the communicator, dataset and variables by default.
-- 
GitLab