{pyactr} Generalize pyactr_print's get_buffer() (#377)

Fixes #376
asmaloney · Aug 13, 2023 · f130a35 · f130a35
1 parent 11e3855
commit f130a35
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/framework/pyactr/pyactr_print.py b/framework/pyactr/pyactr_print.py
@@ -41,10 +41,8 @@ def get_buffer(self, buffer_name: str) -> Buffer:
     """
     Gets a buffer from a name and returns it.
     """
-    if buffer_name == 'goal':
-        return self._ACTRModel__buffers['goal']
-    if buffer_name == 'retrieval':
-        return self._ACTRModel__buffers['retrieval']
+    if buffer_name in self._ACTRModel__buffers:
+        return self._ACTRModel__buffers[buffer_name]
 
     print('ERROR: Buffer \'' + buffer_name + '\' not found')
     raise KeyError