Merge pull request #60 from jeamland/master

Implement forcing of the language from the command line.

1 files changed, 25 insertions, 15 deletions

diff --git a/pycco/main.py b/pycco/main.py
index 487ec5e..0e30ab6 100644
--- a/pycco/main.py
+++ b/pycco/main.py
@@ -33,7 +33,8 @@ Or, to install the latest source
 
 # === Main Documentation Generation Functions ===
 
-def generate_documentation(source, outdir=None, preserve_paths=True):
+def generate_documentation(source, outdir=None, preserve_paths=True,
+                           language=None):
     """
     Generate the documentation for a source file by reading it in, splitting it
     up into comment/code sections, highlighting them for the appropriate
@@ -42,12 +43,13 @@ def generate_documentation(source, outdir=None, preserve_paths=True):
 
     if not outdir:
         raise TypeError("Missing the required 'outdir' keyword argument.")
-    fh = open(source, "r")
-    sections = parse(source, fh.read())
-    highlight(source, sections, preserve_paths=preserve_paths, outdir=outdir)
+    code = open(source, "r").read()
+    language = get_language(source, code, language=language)
+    sections = parse(source, code, language)
+    highlight(source, sections, language, preserve_paths=preserve_paths, outdir=outdir)
     return generate_html(source, sections, preserve_paths=preserve_paths, outdir=outdir)
 
-def parse(source, code):
+def parse(source, code, language):
     """
     Given a string of source code, parse out each comment and the code that
     follows it, and create an individual **section** for it.
@@ -63,7 +65,6 @@ def parse(source, code):
 
     lines = code.split("\n")
     sections = []
-    language = get_language(source)
     has_code = docs_text = code_text = ""
 
     if lines[0].startswith("#!"):
@@ -189,7 +190,7 @@ def preprocess(comment, section_nr, preserve_paths=True, outdir=None):
 
 # === Highlighting the source code ===
 
-def highlight(source, sections, preserve_paths=True, outdir=None):
+def highlight(source, sections, language, preserve_paths=True, outdir=None):
     """
     Highlights a single chunk of code using the **Pygments** module, and runs
     the text of its corresponding comment through **Markdown**.
@@ -201,7 +202,6 @@ def highlight(source, sections, preserve_paths=True, outdir=None):
 
     if not outdir:
         raise TypeError("Missing the required 'outdir' keyword argument.")
-    language = get_language(source)
 
     output = pygments.highlight(language["divider_text"].join(section["code_text"].rstrip() for section in sections),
                                 language["lexer"],
@@ -321,16 +321,20 @@ for ext, l in languages.items():
     # Get the Pygments Lexer for this language.
     l["lexer"] = lexers.get_lexer_by_name(l["name"])
 
-def get_language(source):
+def get_language(source, code, language=None):
     """Get the current language we're documenting, based on the extension."""
 
+    if language is not None:
+        for l in languages.values():
+            if l["name"] == language:
+                return l
+        else:
+            raise ValueError("Unknown forced language: " + language)
+
     m = re.match(r'.*(\..+)', os.path.basename(source))
     if m and m.group(1) in languages:
         return languages[m.group(1)]
     else:
-        source = open(source, "r")
-        code = source.read()
-        source.close()
         lang = lexers.guess_lexer(code).name.lower()
         for l in languages.values():
             if l["name"] == lang:
@@ -387,7 +391,7 @@ highlight_start = "<div class=\"highlight\"><pre>"
 # The end of each Pygments highlight block.
 highlight_end = "</pre></div>"
 
-def process(sources, preserve_paths=True, outdir=None):
+def process(sources, preserve_paths=True, outdir=None, language=None):
     """For each source file passed as argument, generate the documentation."""
 
     if not outdir:
@@ -414,7 +418,8 @@ def process(sources, preserve_paths=True, outdir=None):
                 pass
 
             with open(dest, "w") as f:
-                f.write(generate_documentation(s, preserve_paths=preserve_paths, outdir=outdir))
+                f.write(generate_documentation(s, preserve_paths=preserve_paths, outdir=outdir,
+                                               language=language))
 
             print "pycco = %s -> %s" % (s, dest)
 
@@ -481,9 +486,14 @@ def main():
 
     parser.add_option('-w', '--watch', action='store_true',
                       help='Watch original files and re-generate documentation on changes')
+
+    parser.add_option('-l', '--force-language', action='store', type='string',
+                      dest='language', default=None,
+                      help='Force the language for the given files')
     opts, sources = parser.parse_args()
 
-    process(sources, outdir=opts.outdir, preserve_paths=opts.paths)
+    process(sources, outdir=opts.outdir, preserve_paths=opts.paths,
+            language=opts.language)
 
     # If the -w / --watch option was present, monitor the source directories
     # for changes and re-generate documentation for source files whenever they