Datamodel attributes cache

src/ansys/fluent/core/services/datamodel_se.py

@@ -1,5 +1,6 @@
 """Wrappers over StateEngine based datamodel gRPC service of Fluent."""
 from enum import Enum
+import functools
 import itertools
 import logging
 from typing import Any, Callable, Dict, Iterator, List, Tuple, Type
@@ -274,6 +275,42 @@ def __call__(self, *args, **kwds) -> Any:
         return self.get_state()
 
 
+class EventSubscription:
+    """EventSubscription class for any datamodel event."""
+
+    def __init__(
+        self,
+        service: DatamodelService,
+        request: DataModelProtoModule.SubscribeEventsRequest,
+    ):
+        """Subscribe to a datamodel event."""
+        self._service = service
+        response = service.subscribe_events(request)
+        response = response.response[0]
+        if response.status != DataModelProtoModule.STATUS_SUBSCRIBED:
+            raise RuntimeError(f"Failed to subscribe event: {request}!")
+        self.status = response.status
+        self.tag = response.tag
+        self._service.events[self.tag] = self
+
+    def unsubscribe(self):
+        """Unsubscribe the datamodel event."""
+        if self.status == DataModelProtoModule.STATUS_SUBSCRIBED:
+            self._service.event_streaming.unregister_callback(self.tag)
+            request = DataModelProtoModule.UnsubscribeEventsRequest()
+            request.tag.append(self.tag)
+            response = self._service.unsubscribe_events(request)
+            response = response.response[0]
+            if response.status != DataModelProtoModule.STATUS_UNSUBSCRIBED:
+                raise RuntimeError(f"Failed to unsubscribe event: {request}!")
+            self.status = response.status
+        self._service.events.pop(self.tag, None)
+
+    def __del__(self):
+        """Unsubscribe the datamodel event."""
+        self.unsubscribe()
+
+
 class PyStateContainer(PyCallableStateObject):
     """Object class using StateEngine based DatamodelService as backend. Use
     this class instead of directly calling DatamodelService's method.
@@ -307,8 +344,8 @@ def __init__(self, service: DatamodelService, rules: str, path: Path = None):
             self.path = []
         else:
             self.path = path
-        if pyfluent.DATAMODEL_USE_ATTR_CACHE:
-            self.attrs = {}
+        self.cached_attrs = {}
+        self.event_subscriptions = []
 
     docstring = None
 
@@ -340,6 +377,27 @@ def set_state(self, state: Any = None, **kwargs) -> None:
 
     setState = set_state
 
+    def _get_remote_attr(self, attrib: str) -> Any:
+        request = DataModelProtoModule.GetAttributeValueRequest()
+        request.rules = self.rules
+        request.path = convert_path_to_se_path(self.path)
+        request.attribute = attrib
+        response = self.service.get_attribute_value(request)
+        return _convert_variant_to_value(response.result)
+
+    def _get_cached_attr(self, attrib: str) -> Any:
+        cached_val = self.cached_attrs.get(attrib)
+        if cached_val is None:
+            cached_val = self._get_remote_attr(attrib)
+            self.cached_attrs[attrib] = cached_val
+            self.event_subscriptions.append(
+                self.add_on_attribute_changed(
+                    attrib,
+                    functools.partial(dict.__setitem__, self.cached_attrs, attrib),
+                )
+            )
+        return cached_val
+
     def get_attr(self, attrib: str) -> Any:
         """Get attribute value of the current object.
 
@@ -353,12 +411,9 @@ def get_attr(self, attrib: str) -> Any:
         Any
             Value of the attribute.
         """
-        request = DataModelProtoModule.GetAttributeValueRequest()
-        request.rules = self.rules
-        request.path = convert_path_to_se_path(self.path)
-        request.attribute = attrib
-        response = self.service.get_attribute_value(request)
-        return _convert_variant_to_value(response.result)
+        if pyfluent.DATAMODEL_USE_ATTR_CACHE:
+            return self._get_cached_attr(attrib)
+        return self._get_remote_attr(attrib)
 
     getAttribValue = get_attr
 
@@ -391,41 +446,30 @@ def __call__(self, *args, **kwargs):
 
     docstring = None
 
+    def add_on_attribute_changed(
+        self, attribute: str, cb: Callable
+    ) -> EventSubscription:
+        """Register a callback for when an attribute is changed
 
-class EventSubscription:
-    """EventSubscription class for any datamodel event."""
-
-    def __init__(
-        self,
-        service: DatamodelService,
-        request: DataModelProtoModule.SubscribeEventsRequest,
-    ):
-        """Subscribe to a datamodel event."""
-        self._service = service
-        response = service.subscribe_events(request)
-        response = response.response[0]
-        if response.status != DataModelProtoModule.STATUS_SUBSCRIBED:
-            raise RuntimeError(f"Failed to subscribe event: {request}!")
-        self.status = response.status
-        self.tag = response.tag
-        self._service.events[self.tag] = self
-
-    def unsubscribe(self):
-        """Unsubscribe the datamodel event."""
-        if self.status == DataModelProtoModule.STATUS_SUBSCRIBED:
-            self._service.event_streaming.unregister_callback(self.tag)
-            request = DataModelProtoModule.UnsubscribeEventsRequest()
-            request.tag.append(self.tag)
-            response = self._service.unsubscribe_events(request)
-            response = response.response[0]
-            if response.status != DataModelProtoModule.STATUS_UNSUBSCRIBED:
-                raise RuntimeError(f"Failed to unsubscribe event: {request}!")
-            self.status = response.status
-        self._service.events.pop(self.tag, None)
+        Parameters
+        ----------
+        attribute : str
+            attribute name
+        cb : Callable
+            Callback function
 
-    def __del__(self):
-        """Unsubscribe the datamodel event."""
-        self.unsubscribe()
+        Returns
+        -------
+        EventSubscription
+            EventSubscription instance which can be used to unregister the callback
+        """
+        request = DataModelProtoModule.SubscribeEventsRequest()
+        e = request.eventrequest.add(rules=self.rules)
+        e.attributeChangedEventRequest.path = convert_path_to_se_path(self.path)
+        e.attributeChangedEventRequest.attribute = attribute
+        subscription = EventSubscription(self.service, request)
+        self.service.event_streaming.register_callback(subscription.tag, self, cb)
+        return subscription
 
 
 class PyMenu(PyStateContainer):
@@ -607,31 +651,6 @@ def add_on_affected_at_type_path(
         self.service.event_streaming.register_callback(subscription.tag, self, cb)
         return subscription
 
-    def add_on_attribute_changed(
-        self, attribute: str, cb: Callable
-    ) -> EventSubscription:
-        """Register a callback for when an attribute is changed
-
-        Parameters
-        ----------
-        attribute : str
-            attribute name
-        cb : Callable
-            Callback function
-
-        Returns
-        -------
-        EventSubscription
-            EventSubscription instance which can be used to unregister the callback
-        """
-        request = DataModelProtoModule.SubscribeEventsRequest()
-        e = request.eventrequest.add(rules=self.rules)
-        e.attributeChangedEventRequest.path = convert_path_to_se_path(self.path)
-        e.attributeChangedEventRequest.attribute = attribute
-        subscription = EventSubscription(self.service, request)
-        self.service.event_streaming.register_callback(subscription.tag, self, cb)
-        return subscription
-
     def add_on_command_attribute_changed(
         self, command: str, attribute: str, cb: Callable
     ) -> EventSubscription:
@@ -715,31 +734,6 @@ def add_on_changed(self, cb: Callable) -> EventSubscription:
         self.service.event_streaming.register_callback(subscription.tag, self, cb)
         return subscription
 
-    def add_on_attribute_changed(
-        self, attribute: str, cb: Callable
-    ) -> EventSubscription:
-        """Register a callback for when an attribute is changed
-
-        Parameters
-        ----------
-        attribute : str
-            attribute name
-        cb : Callable
-            Callback function
-
-        Returns
-        -------
-        EventSubscription
-            EventSubscription instance which can be used to unregister the callback
-        """
-        request = DataModelProtoModule.SubscribeEventsRequest()
-        e = request.eventrequest.add(rules=self.rules)
-        e.attributeChangedEventRequest.path = convert_path_to_se_path(self.path)
-        e.attributeChangedEventRequest.attribute = attribute
-        subscription = EventSubscription(self.service, request)
-        self.service.event_streaming.register_callback(subscription.tag, self, cb)
-        return subscription
-
 
 def _bool_value_if_none(val, default):
     if isinstance(val, bool) or val is None: