Cleanups

1 files changed, 5 insertions, 4 deletions

diff --git a/prolog/util.py b/prolog/util.py
index b0447a6..fc243f2 100644
--- a/prolog/util.py
+++ b/prolog/util.py
@@ -36,13 +36,14 @@ def stringify(tokens, indent=''):
         else:
             s += t[1]
 
-    return s
+    return s.strip().replace('\n', ' ').replace('\t', ' ')
 
 # return a list of lines in 'code', and a list of rule indexes
 def decompose(code):
     lines = []
     rules = []
     tokens = tokenize(code)
+    tokens.append(('EOF', ''))
 
     line = []
     parens = []
@@ -54,11 +55,11 @@ def decompose(code):
             line = []
             continue
         if not parens:
-            if t[0] in ('PERIOD', 'FROM', 'COMMA'):
+            if t[0] in ('PERIOD', 'FROM', 'COMMA', 'EOF'):
                 if line != []:
                     lines.append(line[:])
                     line = []
-                if t[0] == 'PERIOD':
+                if t[0] in ('PERIOD', 'EOF'):
                     rules.append((rule_start, len(lines)))
                     rule_start = len(lines)
                 continue
@@ -91,7 +92,7 @@ def compose(lines, rules):
                 if line and line[-1][0] != 'SEMI' and i < end-1 and lines[i+1][-1][0] != 'SEMI':
                     code += ','
                 code += '\n'
-    return code
+    return code.strip()
 
 # standardize variable names in order of appearance
 def rename_vars(tokens, names={}):