Schema.has_entity and Schema.has_field

westernx · Oct 29, 2015 · 94f6c7b · 94f6c7b
1 parent b34eeb4
commit 94f6c7b
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 1 deletion.
diff --git a/sgschema/schema.py b/sgschema/schema.py
@@ -414,6 +414,28 @@ def resolve_one_field(self, entity_type, field_spec, **kwargs):
         else:
             raise ValueError('%r returned %s %s fields' % (field_spec, len(res), entity_type))
 
+    def has_entity(self, entity_spec, **kwargs):
+        kwargs['strict'] = True
+        try:
+            entities = self.resolve_entity(entity_spec, **kwargs)
+            return bool(entities)
+        except ValueError as e:
+            if 'is not an entity' in e.args[0]:
+                return False
+            else:
+                raise
+
+    def has_field(self, entity_type, field_spec, **kwargs):
+        kwargs['strict'] = True
+        try:
+            fields = self.resolve_field(entity_type, field_spec, **kwargs)
+            return bool(fields)
+        except ValueError as e:
+            if 'is not a field' in e.args[0]:
+                return False
+            else:
+                raise
+
     def resolve_structure(self, x, entity_type=None, **kwargs):
         """Traverse a nested structure resolving names in entities.
 

diff --git a/tests/test_entities.py b/tests/test_entities.py
@@ -46,4 +46,12 @@ def test_one(self):
         self.assertEqual(self.s.resolve_one_entity('$Alias'), 'Entity')
         self.assertEqual(self.s.resolve_one_entity('#TagOne'), 'Entity')
         self.assertRaises(ValueError, self.s.resolve_one_entity, '#TagNone')
-        self.assertRaises(ValueError, self.s.resolve_one_entity, '#TagTwo')
+        self.assertRaises(ValueError, self.s.resolve_one_entity, '#TagTwo')
+
+    def test_has_entity(self):
+        self.assertTrue(self.s.has_entity('Entity'))
+        self.assertTrue(self.s.has_entity('!Entity'))
+        self.assertTrue(self.s.has_entity('$Alias'))
+        self.assertTrue(self.s.has_entity('#TagOne'))
+        self.assertFalse(self.s.has_entity('NotAnEntity'))
+        self.assertFalse(self.s.has_entity('#TagNone'))
diff --git a/tests/test_fields.py b/tests/test_fields.py
@@ -76,3 +76,13 @@ def test_many(self):
             'sg_type', 'sg_version', 'attr', 'multi_a', 'multi_b',
         ])
 
+    def test_has_field(self):
+        self.assertTrue(self.s.has_field('Entity', 'sg_type'))
+        self.assertTrue(self.s.has_field('Entity', 'version'))
+        self.assertTrue(self.s.has_field('Entity', '$alias'))
+        self.assertTrue(self.s.has_field('Entity', '#tagone'))
+        self.assertTrue(self.s.has_field('Entity', '#tagtwo'))
+        self.assertFalse(self.s.has_field('Entity', '#tagnone'))
+        self.assertFalse(self.s.has_field('Entity', 'notafield'))
+        self.assertRaises(ValueError, self.s.has_field, 'NotAnEntity', 'whatever')
+