From a1ece67b63f3a91084fd19a68801e85f8c2a7692 Mon Sep 17 00:00:00 2001
From: "Ankur Sinha (Ankur Sinha Gmail)" <sanjay.ankur@gmail.com>
Date: Wed, 11 Sep 2024 17:11:05 +0100
Subject: [PATCH] feat(nml): use cached property values

Instead of looking up again and caching independently
---
 neuroml/nml/helper_methods.py | 12 +++++-------
 neuroml/nml/nml.py            | 21 +++++++++++----------
 neuroml/test/test_cell.py     |  2 +-
 3 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/neuroml/nml/helper_methods.py b/neuroml/nml/helper_methods.py
index 0c23a0e..08c07f7 100644
--- a/neuroml/nml/helper_methods.py
+++ b/neuroml/nml/helper_methods.py
@@ -911,21 +911,19 @@ def __str__(self):
     source='''\
 
     # Get segment object by its id
-    @lru_cache(maxsize=1000)
     def get_segment(self, segment_id: int) -> Segment:
         """Get segment object by its id
 
         :param segment_id: ID of segment
+        :type segment_id: int
         :return: segment
 
         :raises ValueError: if the segment is not found in the cell
         """
-
-        for segment in self.morphology.segments:
-            if segment.id == segment_id:
-                return segment
-
-        raise ValueError("Segment with id "+str(segment_id)+" not found in cell "+str(self.id))
+        try:
+            return self.segment_ids_vs_segments[segment_id]
+        except KeyError:
+            raise ValueError("Segment with id "+str(segment_id)+" not found in cell "+str(self.id))
 
     def get_segments_by_substring(self, substring: str) -> typing.Dict[str, Segment]:
         """Get a dictionary of segment IDs and the segment matching the specified substring
diff --git a/neuroml/nml/nml.py b/neuroml/nml/nml.py
index 3cb2d7d..8e08329 100644
--- a/neuroml/nml/nml.py
+++ b/neuroml/nml/nml.py
@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 
 #
-# Generated Wed Sep 11 16:48:39 2024 by generateDS.py version 2.44.1.
+# Generated Wed Sep 11 17:09:38 2024 by generateDS.py version 2.44.1.
 # Python 3.11.9 (main, Aug 23 2024, 00:00:00) [GCC 14.2.1 20240801 (Red Hat 14.2.1-1)]
 #
 # Command line options:
@@ -48430,23 +48430,24 @@ def _buildChildren(
         super(Cell, self)._buildChildren(child_, node, nodeName_, True)
 
     # Get segment object by its id
-    @lru_cache(maxsize=1000)
     def get_segment(self, segment_id: int) -> Segment:
         """Get segment object by its id
 
         :param segment_id: ID of segment
+        :type segment_id: int
         :return: segment
 
         :raises ValueError: if the segment is not found in the cell
         """
-
-        for segment in self.morphology.segments:
-            if segment.id == segment_id:
-                return segment
-
-        raise ValueError(
-            "Segment with id " + str(segment_id) + " not found in cell " + str(self.id)
-        )
+        try:
+            return self.segment_ids_vs_segments[segment_id]
+        except KeyError:
+            raise ValueError(
+                "Segment with id "
+                + str(segment_id)
+                + " not found in cell "
+                + str(self.id)
+            )
 
     def get_segments_by_substring(self, substring: str) -> typing.Dict[str, Segment]:
         """Get a dictionary of segment IDs and the segment matching the specified substring
diff --git a/neuroml/test/test_cell.py b/neuroml/test/test_cell.py
index 16c883e..6e6e24e 100644
--- a/neuroml/test/test_cell.py
+++ b/neuroml/test/test_cell.py
@@ -395,7 +395,7 @@ def test_get_morphology_root(self):
 
             # clear the cache, otherwise it'll return the old value again in
             # other function calls also
-            acell.get_segment.cache_clear()
+            acell.get_segment_ids_vs_segments(cached=False)
 
             # also update all descendents to ensure cell remains valid
             for seg in acell.morphology.segments: