Abstract out write_genotype.

* pyhegp/serialization.py (write_genotype): New function.
* pyhegp/pyhegp.py: Import write_genotype from pyhegp.serialization.
(encrypt, cat): Use write_genotype.

2 files changed, 7 insertions, 6 deletions

diff --git a/pyhegp/pyhegp.py b/pyhegp/pyhegp.py
index 3229724..8de6f9d 100644
--- a/pyhegp/pyhegp.py
+++ b/pyhegp/pyhegp.py
@@ -22,7 +22,7 @@ import click
 import numpy as np
 from scipy.stats import special_ortho_group
 
-from pyhegp.serialization import Summary, read_summary, write_summary, read_genotype
+from pyhegp.serialization import Summary, read_summary, write_summary, read_genotype, write_genotype
 
 Stats = namedtuple("Stats", "n mean std")
 
@@ -108,7 +108,7 @@ def encrypt(genotype_file, summary_file, key_file, ciphertext_file):
                                       key)
     if key_file:
         np.savetxt(key_file, key, delimiter=",", fmt="%f")
-    np.savetxt(ciphertext_file, encrypted_genotype, delimiter=",", fmt="%f")
+    write_genotype(ciphertext_file, encrypted_genotype)
 
 @main.command()
 @click.option("--output", "-o", "output_file",
@@ -117,10 +117,8 @@ def encrypt(genotype_file, summary_file, key_file, ciphertext_file):
               help="output file")
 @click.argument("ciphertext-files", type=click.File("rb"), nargs=-1)
 def cat(output_file, ciphertext_files):
-    np.savetxt(output_file,
-               np.vstack([read_genotype(file) for file in ciphertext_files]),
-               delimiter=",",
-               fmt="%f")
+    write_genotype(output_file,
+                   np.vstack([read_genotype(file) for file in ciphertext_files]))
 
 if __name__ == "__main__":
     main()
diff --git a/pyhegp/serialization.py b/pyhegp/serialization.py
index 9b0401a..269528e 100644
--- a/pyhegp/serialization.py
+++ b/pyhegp/serialization.py
@@ -54,3 +54,6 @@ def write_summary(file, summary):
 
 def read_genotype(genotype_file):
     return np.loadtxt(genotype_file, delimiter=",")
+
+def write_genotype(file, genotype):
+    np.savetxt(file, genotype, delimiter=",", fmt="%f")