add term_names property; closes #20

bambinos · Aug 27, 2016 · aa535c5 · aa535c5
1 parent 6d56b50
commit aa535c5
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/bambi/models.py b/bambi/models.py
@@ -417,6 +417,10 @@ def add_term(self, variable, data=None, label=None, categorical=False,
         self.terms[term.name] = term
         self.built = False
 
+    @property
+    def term_names(self):
+        return list(self.terms.keys())
+
 
 class Term(object):
 

diff --git a/bambi/tests/test_model.py b/bambi/tests/test_model.py
@@ -79,6 +79,14 @@ def test_model_init_and_intercept(diabetes_data):
     assert not model.terms
 
 
+def test_model_term_names_property(diabetes_data):
+    model = Model(diabetes_data)
+    model.add_term('BMI')
+    model.add_term('BP')
+    model.add_term('S1')
+    assert model.term_names == ['BMI', 'BP', 'S1']
+
+
 def test_add_term_to_model(base_model):
 
     base_model.add_term('BMI')