From 1d90b5a4677d6310ed4770fd1faad5e28ef900e9 Mon Sep 17 00:00:00 2001
From: Thomas Colthurst <thomaswc@google.com>
Date: Tue, 24 Sep 2024 18:07:40 +0000
Subject: [PATCH] Fix load_schema bugs and add test

---
 cxx/assets/animals.unary.with_comments.schema | 86 +++++++++++++++++++
 py/hirm_io.py                                 |  4 +-
 py/hirm_io_test.py                            | 10 +++
 3 files changed, 98 insertions(+), 2 deletions(-)
 create mode 100644 cxx/assets/animals.unary.with_comments.schema

diff --git a/cxx/assets/animals.unary.with_comments.schema b/cxx/assets/animals.unary.with_comments.schema
new file mode 100644
index 0000000..aa86db8
--- /dev/null
+++ b/cxx/assets/animals.unary.with_comments.schema
@@ -0,0 +1,86 @@
+# Schema for animal attributes expressed as unary relations.
+black ~ bernoulli(animal)
+white ~ bernoulli(animal)
+blue ~ bernoulli(animal)
+brown ~ bernoulli(animal)
+gray ~ bernoulli(animal)
+orange ~ bernoulli(animal)
+red ~ bernoulli(animal)
+yellow ~ bernoulli(animal)
+patches ~ bernoulli(animal)
+spots ~ bernoulli(animal)
+stripes ~ bernoulli(animal)
+furry ~ bernoulli(animal)
+hairless ~ bernoulli(animal)
+toughskin ~ bernoulli(animal)
+big ~ bernoulli(animal)  # TODO(thomaswc): Add definitions of ambiguous properties
+small ~ bernoulli(animal)
+bulbous ~ bernoulli(animal)
+lean ~ bernoulli(animal)
+flippers ~ bernoulli(animal)
+hands ~ bernoulli(animal)
+hooves ~ bernoulli(animal)
+pads ~ bernoulli(animal)
+paws ~ bernoulli(animal)
+longleg ~ bernoulli(animal)
+longneck ~ bernoulli(animal)
+tail ~ bernoulli(animal)
+chewteeth ~ bernoulli(animal)
+meatteeth ~ bernoulli(animal)
+buckteeth ~ bernoulli(animal)
+strainteeth ~ bernoulli(animal)
+horns ~ bernoulli(animal)
+claws ~ bernoulli(animal)
+tusks ~ bernoulli(animal)
+smelly ~ bernoulli(animal)
+flys ~ bernoulli(animal)
+hops ~ bernoulli(animal)
+swims ~ bernoulli(animal)
+tunnels ~ bernoulli(animal)
+walks ~ bernoulli(animal)
+fast ~ bernoulli(animal)
+slow ~ bernoulli(animal)
+strong ~ bernoulli(animal)
+weak ~ bernoulli(animal)
+muscle ~ bernoulli(animal)
+bipedal ~ bernoulli(animal)
+quadrapedal ~ bernoulli(animal)
+active ~ bernoulli(animal)
+inactive ~ bernoulli(animal)
+nocturnal ~ bernoulli(animal)
+hibernate ~ bernoulli(animal)
+agility ~ bernoulli(animal)
+fish ~ bernoulli(animal)  # Is this the verb or the noun?
+meat ~ bernoulli(animal)
+plankton ~ bernoulli(animal)
+vegetation ~ bernoulli(animal)
+insects ~ bernoulli(animal)
+forager ~ bernoulli(animal)
+grazer ~ bernoulli(animal)
+hunter ~ bernoulli(animal)
+scavenger ~ bernoulli(animal)
+skimmer ~ bernoulli(animal)
+stalker ~ bernoulli(animal)
+newworld ~ bernoulli(animal)
+oldworld ~ bernoulli(animal)
+arctic ~ bernoulli(animal)
+coastal ~ bernoulli(animal)
+desert ~ bernoulli(animal)
+bush ~ bernoulli(animal)
+plains ~ bernoulli(animal)
+forest ~ bernoulli(animal)
+fields ~ bernoulli(animal)
+jungle ~ bernoulli(animal)
+mountains ~ bernoulli(animal)
+ocean ~ bernoulli(animal)
+ground ~ bernoulli(animal)
+water ~ bernoulli(animal)
+tree ~ bernoulli(animal)
+cave ~ bernoulli(animal)
+fierce ~ bernoulli(animal)
+timid ~ bernoulli(animal)
+smart ~ bernoulli(animal)
+group ~ bernoulli(animal)
+solitary ~ bernoulli(animal)
+nestspot ~ bernoulli(animal)
+domestic ~ bernoulli(animal)
diff --git a/py/hirm_io.py b/py/hirm_io.py
index 22b1931..87a3a0f 100644
--- a/py/hirm_io.py
+++ b/py/hirm_io.py
@@ -23,11 +23,11 @@ def load_schema(path):
   relations = {}
   comment_line_re = re.compile(r'\s*#.*')
   # TODO(thomaswc): Handle distribution parameters in brackets
-  line_re = re.compile(r'\s*(\w+)\s*~\s*(\w+)\s*\(\s*(.+)\s*\)\s*(#.*)?')
+  line_re = re.compile(r'\s*(\w+)\s*~\s*(\w+)\s*\(\s*([^\)]+)\s*\)\s*(#.*)?')
   with open(path, 'r') as f:
     for line in f:
       if comment_line_re.match(line):
-        next
+        continue
       m = line_re.match(line)
       if not m:
         print(f"Could not parse schema line\n{line}\n")
diff --git a/py/hirm_io_test.py b/py/hirm_io_test.py
index 73e3e42..2bedae5 100755
--- a/py/hirm_io_test.py
+++ b/py/hirm_io_test.py
@@ -12,6 +12,16 @@ def test_load_schema(self):
     self.assertEqual(relations["black"].distribution, "bernoulli")
     self.assertEqual(relations["black"].domains, ["animal"])
 
+  def test_load_schema_with_comments(self):
+    relations = hirm_io.load_schema("../cxx/assets/animals.unary.with_comments.schema")
+    self.assertEqual(len(relations), 85)
+    self.assertEqual(relations["black"].name, "black")
+    self.assertEqual(relations["black"].distribution, "bernoulli")
+    self.assertEqual(relations["black"].domains, ["animal"])
+    self.assertEqual(relations["big"].name, "big")
+    self.assertEqual(relations["big"].distribution, "bernoulli")
+    self.assertEqual(relations["big"].domains, ["animal"])
+
   def test_load_observations(self):
     observations = hirm_io.load_observations("../cxx/assets/animals.unary.obs")
     self.assertEqual(len(observations), 4250)