From 28381d12fd9de54c0d69e2df76baf29e75ade6b5 Mon Sep 17 00:00:00 2001
From: Jameson Graef Rollins <jameson.rollins@ligo.org>
Date: Mon, 22 Feb 2021 10:52:59 -0800
Subject: [PATCH] add ability to retrieve nested items from Budget and Trace
objects
e.g. Budget['SubBuget.BudgetItem']
---
gwinc/nb.py | 24 ++++++++++++++++++++++--
gwinc/trace.py | 19 +++++++++++++++++--
2 files changed, 39 insertions(+), 4 deletions(-)
diff --git a/gwinc/nb.py b/gwinc/nb.py
index 19ce472e..56baa8bf 100644
--- a/gwinc/nb.py
+++ b/gwinc/nb.py
@@ -418,13 +418,33 @@ class Budget(Noise):
return name
def __getitem__(self, name):
+ """Get a (possibly nested) sub-BudgetItem.
+
+ """
+ try:
+ name, rest = name.split('.', 1)
+ except ValueError:
+ rest = None
try:
- return self._noise_objs[name]
+ o = self._noise_objs[name]
except KeyError:
try:
- return self._cal_objs[name]
+ o = self._cal_objs[name]
except KeyError:
raise KeyError("unknown noise or cal name '{}".format(name))
+ if rest:
+ return o[rest]
+ else:
+ return o
+
+ def get(self, key, default=None):
+ """Get a (possibly nested) sub-BudgetItem.
+
+ """
+ try:
+ return self[key]
+ except KeyError:
+ return default
def keys(self):
"""Iterate over budget noise names."""
diff --git a/gwinc/trace.py b/gwinc/trace.py
index d113b985..a15d5590 100644
--- a/gwinc/trace.py
+++ b/gwinc/trace.py
@@ -61,8 +61,23 @@ class BudgetTrace:
raise AttributeError
def __getitem__(self, name):
- """get budget trace by name"""
- return self._bdict[name]
+ """get budget trace by name
+
+ """
+ try:
+ name, rest = name.split('.', 1)
+ return self._bdict[name][rest]
+ except ValueError:
+ return self._bdict[name]
+
+ def get(self, key, default=None):
+ """get a (possibly nested) Trace item.
+
+ """
+ try:
+ return self[key]
+ except KeyError:
+ return default
def items(self):
"""iterator of budget (name, trace) tuples"""
--
GitLab