namegen.py: use json instead of yaml to save stats.

yaml is slow and json is faster and available in the standard library.

1 files changed, 9 insertions, 17 deletions

diff --git a/namegen.py b/namegen.py
index 194208f..baed954 100755
--- a/namegen.py
+++ b/namegen.py
@@ -43,16 +43,10 @@ LICENSE
 
 """
 
+import json
 import sys
 import random
 
-try:
-    import yaml
-    skip_yaml = False
-except ImportError:
-    sys.stderr.write("Could not find module YAML: will not be able to save / load statistics\n")
-    skip_yaml = True
-
 class Stats:
 
     def __init__(self,context):
@@ -73,6 +67,7 @@ class Stats:
         """Calculate stats from a list of names in a file-like object."""
 
         for name in fp.readlines():
+            name = name.decode('utf-8')
             for i in xrange(min(self.context,len(name))):
                 self._increment_stats(name[0:i],name[i])
             for i in xrange(self.context,len(name)):
@@ -80,17 +75,11 @@ class Stats:
 
     def load_stats(self,fp):
         """Load stats from a yaml in a string or file-like object"""
-        if skip_yaml:
-            return False
-        self.stats = yaml.safe_load(fp)
+        self.stats = json.load(fp)
 
     def save_stats(self,fp):
         """Save stats in yaml to a file-like object"""
-        if skip_yaml:
-            return False
-        yaml.safe_dump(self.stats,fp)
-        return True
-
+        json.dump(self.stats,fp)
 
     def generate_name(self):
         """Generate a name following the stats"""
@@ -110,8 +99,11 @@ class Stats:
 
 def main():
     mystats = Stats(3)
-    fp = open('data/nomif2.txt')
-    mystats.read_names(fp)
+    #fp = open('data/nomif2.txt')
+    #mystats.read_names(fp)
+    #fp.close()
+    fp = open('it_female.stats')
+    mystats.load_stats(fp)
     fp.close()
     print mystats.generate_name(),