Fix auto modifier

sbrunner · Dec 20, 2013 · cc047e4 · cc047e4
1 parent f5f3845
commit cc047e4
Showing 1 changed file with 20 additions and 7 deletions.
diff --git a/sphinx-prompt/__init__.py b/sphinx-prompt/__init__.py
@@ -15,15 +15,15 @@
 
 class PromptDirective(rst.Directive):
 
-    optional_arguments = 2
+    optional_arguments = 3
     has_content = True
 
     def run(self):
         self.assert_has_content()
 
         language = 'text'
         prompt = None
-        modifier = []
+        modifiers = []
 
         if self.arguments:
             language = self.arguments[0]
@@ -32,13 +32,13 @@ def run(self):
             elif language == 'bash':
                 prompt = '$'
             if len(self.arguments) > 2:
-                modifier = self.arguments[2].split(',')
-            if modifier == 'auto':
+                modifiers = self.arguments[2].split(',')
+            if 'auto' in modifiers:
                 prompts = prompt.split(',')
 
         html = '<pre class="highlight">'
         html += '<style type="text/css" scoped>'
-        if modifier == 'auto':
+        if 'auto' in modifiers:
             for index, prompt in enumerate(prompts):
                 html += """span.prompt%i:before {
   content: "%s ";
@@ -64,8 +64,9 @@ def run(self):
         statement = []
         prompt_index = -1
         for line in self.content:
-            if modifier == 'auto':
+            if 'auto' in modifiers:
                 latex += '\n' + line
+
                 for index, prompt in enumerate(prompts):
                     if line.startswith(prompt):
                         if len(statement) > 0:
@@ -81,10 +82,11 @@ def run(self):
                         line = line[len(prompt):].strip()
                         prompt_index = index
                         break
+
                 statement.append(line)
             elif language == 'bash' or language == 'python':
                 statement.append(line)
-                if not line[-1] == '\\':
+                if len(line) == 0 or not line[-1] == '\\':
                     html += '<span class="prompt">%s</span>\n' % (
                         highlight(
                             '\n'.join(statement),
@@ -110,6 +112,17 @@ def run(self):
                 else:
                     latex += '\n' + line
 
+        # Add last prompt
+        if 'auto' in modifiers and len(statement) > 0:
+            html += '<span class="prompt%i">%s</span>\n' % (
+                prompt_index,
+                highlight(
+                    '\n'.join(statement),
+                    Lexer(),
+                    HtmlFormatter(nowrap=True)
+                ).strip('\r\n')
+            )
+
         html += "</pre>"
         latex += "\n\\end{Verbatim}"