FormFieldClassifier.form_classes and FormFieldClassifier.field_classes

TeamHG-Memex · Feb 9, 2016 · e5a7b09 · e5a7b09
1 parent 5b5086c
commit e5a7b09
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/formasaurus/classifiers.py b/formasaurus/classifiers.py
@@ -227,6 +227,16 @@ def _cached_model_path(cls):
         path = "formasaurus-%s.joblib" % dependencies_string()
         return at_root(path)
 
+    @property
+    def form_classes(self):
+        """ Possible form classes """
+        return self.form_classifier.classes
+
+    @property
+    def field_classes(self):
+        """ Possible field classes """
+        return self._field_model.classes_
+
 
 class FormClassifier(object):
     """

diff --git a/tests/test_classifiers.py b/tests/test_classifiers.py
@@ -29,6 +29,12 @@ def test_extract_forms_no_fields_direct(tree):
     assert forms[0][1] == 'login'
 
 
+def test_classes():
+    ex = classifiers.get_instance()
+    assert 'registration' in ex.form_classes
+    assert 'password' in ex.field_classes
+
+
 @pytest.mark.parametrize(['fields'], [[True], [False]])
 def test_extract_forms_proba(tree, fields):
     forms = formasaurus.extract_forms(tree, proba=True, threshold=0, fields=fields)