Write to a stream instead of file name in exporters that support it.

4 files changed, 41 insertions, 22 deletions

diff --git a/screenplain/export/annotated_html.py b/screenplain/export/annotated_html.py
index 2bcf643..64c38ee 100644
--- a/screenplain/export/annotated_html.py
+++ b/screenplain/export/annotated_html.py
@@ -1,11 +1,11 @@
 import sys
 import re
 import cgi
-from screenplain.parse import parse
-import screenplain.parse
+from screenplain.types import *
 
 def unspace(text):
-    text = re.sub(r'\s+', ' ', text)
+    text = re.sub(r'\s*\n\s*', '\n', text)
+    text = re.sub(r'\s\s+', ' ', text)
     text = re.sub(r'>\s+<', '><', text)
     return text.strip()
 
@@ -18,19 +18,18 @@ paragraph_html = unspace("""
 """)
 
 types = {
-    screenplain.parse.Slug: 'Slug',
-    screenplain.parse.Dialog: 'Dialog',
-    screenplain.parse.DualDialog: 'Dual',
-    screenplain.parse.Action: 'Action',
-    screenplain.parse.Transition: 'Transition',
+    Slug: 'Slug',
+    Dialog: 'Dialog',
+    DualDialog: 'Dual',
+    Action: 'Action',
+    Transition: 'Transition',
 }
 
 def to_html(text):
     return re.sub('  ', '&nbsp; ', cgi.escape(text))
 
-def to_annotated_html(input, out):
-    paragraphs = parse(input)
-    for para in paragraphs:
+def to_annotated_html(screenplay, out):
+    for para in screenplay:
         lines = para.format()
         margin = '<p>&nbsp;</p>' * para.top_margin
         html_text = ''.join(
diff --git a/screenplain/export/pdf.py b/screenplain/export/pdf.py
index f7c9f24..d1b920b 100644
--- a/screenplain/export/pdf.py
+++ b/screenplain/export/pdf.py
@@ -11,8 +11,9 @@ def to_pdf(screenplay, output_file):
     c = canvas.Canvas(output_file, pagesize=pagesizes.A4)
     c.setFont('Courier', 12)
 
-    for page_no, page in enumerate(get_pages(screenplay), 1):
-        if page_no != 1:
+    for page_no, page in enumerate(get_pages(screenplay)):
+        # page numbers are 0-based
+        if page_no != 0:
             c.showPage()
         c.setFont('Courier', 12)
         for line_no, line in enumerate(page):
diff --git a/screenplain/export/text.py b/screenplain/export/text.py
index 12b8723..d4c5b77 100644
--- a/screenplain/export/text.py
+++ b/screenplain/export/text.py
@@ -2,8 +2,7 @@ import sys
 import codecs
 from screenplain.format import get_pages
 
-def to_text(screenplay, output_file):
-    out = codecs.open(output_file, 'w', 'utf-8')
+def to_text(screenplay, out):
     for page_no, page in enumerate(get_pages(screenplay)):
         # page_no is 0-based
         if page_no != 0:
diff --git a/screenplain/main.py b/screenplain/main.py
index 469da0f..a2d6c1f 100644
--- a/screenplain/main.py
+++ b/screenplain/main.py
@@ -10,16 +10,21 @@ from optparse import OptionParser
 
 from screenplain.parsers.spmd import parse
 
+output_formats = (
+    'text', 'pdf', 'annotated_html'
+)
+
 usage = 'Usage: %prog [options] input-file output-file'
 
-if __name__ == '__main__':
+def main(args):
     parser = OptionParser(usage=usage)
     parser.add_option(
         '-f', '--format', dest='output_format',
         metavar='FORMAT',
-        help='Set what kind of file to create. FORMAT can be pdf or text.'
+        help=('Set what kind of file to create. FORMAT can be any of '
+            + ', '.join(output_formats))
     )
-    options, args = parser.parse_args()
+    options, args = parser.parse_args(args)
     if len(args) != 2:
         parser.error('Expected input-file and output-file arguments')
     input_file, output_file = args
@@ -29,13 +34,28 @@ if __name__ == '__main__':
             options.output_format = 'pdf'
         else:
             options.output_format = 'text'
+    if options.output_format not in output_formats:
+        parser.error(
+            'Unknown output format. Expected one of: ' +
+            ', '.join(output_formats))
 
     input = codecs.open(input_file, 'r', 'utf-8')
     screenplay = parse(input)
 
-    if options.output_format == 'text':
-        from screenplain.export.text import to_text
-        to_text(screenplay, output_file)
-    else:
+    if options.output_format == 'pdf':
         from screenplain.export.pdf import to_pdf
         to_pdf(screenplay, output_file)
+    else:
+        output = codecs.open(output_file, 'w', 'utf-8')
+        try:
+            if options.output_format == 'text':
+                from screenplain.export.text import to_text
+                to_text(screenplay, output)
+            elif options.output_format == 'annotated_html':
+                from screenplain.export.annotated_html import to_annotated_html
+                to_annotated_html(screenplay, output)
+        finally:
+            output.close()
+
+if __name__ == '__main__':
+    main(sys.argv[1:])