Basic Python 3 support

4 files changed, 51 insertions, 28 deletions

diff --git a/.travis.yml b/.travis.yml
index bfbe563..62e7c6b 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -1,10 +1,12 @@
 language: python
 python:
     - '2.7'
+    - '3.5'
 install:
     - 'pip install -r requirements.txt'
     - 'pip install -r requirements.test.txt'
 script:
     - 'py.test --cov=pycco tests/'
+    - 'python -m pycco.main pycco/main.py'
 after_success:
     - coveralls
diff --git a/pycco/compat.py b/pycco/compat.py
new file mode 100644
index 0000000..6660531
--- /dev/null
+++ b/pycco/compat.py
@@ -0,0 +1,4 @@
+try:
+    pycco_unichr = unichr
+except NameError:
+    pycco_unichr = chr
diff --git a/pycco/main.py b/pycco/main.py
index df2b2bc..cde05d7 100644
--- a/pycco/main.py
+++ b/pycco/main.py
@@ -1,4 +1,5 @@
 #!/usr/bin/env python
+from __future__ import print_function
 
 """
 "**Pycco**" is a Python port of [Docco](http://jashkenas.github.com/docco/):
@@ -35,7 +36,7 @@ Or, to install the latest source
 
 
 def generate_documentation(source, outdir=None, preserve_paths=True,
-                           language=None):
+                           language=None, encoding="utf8"):
     """
     Generate the documentation for a source file by reading it in, splitting it
     up into comment/code sections, highlighting them for the appropriate
@@ -44,7 +45,7 @@ def generate_documentation(source, outdir=None, preserve_paths=True,
 
     if not outdir:
         raise TypeError("Missing the required 'outdir' keyword argument.")
-    code = open(source, "r").read()
+    code = open(source, "rb").read().decode(encoding)
     return _generate_documentation(source, code, outdir, preserve_paths, language)
 
 
@@ -226,6 +227,8 @@ def highlight(sections, language, preserve_paths=True, outdir=None):
             docs_text = unicode(section["docs_text"])
         except UnicodeError:
             docs_text = unicode(section["docs_text"].decode('utf-8'))
+        except NameError:
+            docs_text = section['docs_text']
         section["docs_html"] = markdown(preprocess(docs_text,
                                                    preserve_paths=preserve_paths,
                                                    outdir=outdir))
@@ -361,9 +364,9 @@ def get_language(source, code, language=None):
             else:
                 raise ValueError()
         except ValueError:
-                # If pygments can't find any lexers, it will raise its own
-                # subclass of ValueError. We will catch it and raise ours
-                # for consistency.
+            # If pygments can't find any lexers, it will raise its own
+            # subclass of ValueError. We will catch it and raise ours
+            # for consistency.
             raise ValueError("Can't figure out the language!")
 
 
@@ -403,15 +406,20 @@ def shift(list, default):
         return default
 
 
+def remove_control_chars(s):
+    # Sanitization regexp copied from
+    # http://stackoverflow.com/questions/92438/stripping-non-printable-characters-from-a-string-in-python
+    from pycco.compat import pycco_unichr
+    control_chars = ''.join(map(pycco_unichr, list(range(0, 32)) + list(range(127, 160))))
+    control_char_re = re.compile(u'[{}]'.format(re.escape(control_chars)))
+    return control_char_re.sub('', s)
+
+
 def ensure_directory(directory):
     """
     Sanitize directory string and ensure that the destination directory exists.
     """
-    # Sanitization regexp copied from
-    # http://stackoverflow.com/questions/92438/stripping-non-printable-characters-from-a-string-in-python
-    control_chars = ''.join(map(unichr, range(0, 32) + range(127, 160)))
-    control_char_re = re.compile(u'[{}]'.format(re.escape(control_chars)))
-    directory = control_char_re.sub('', directory)
+    directory = remove_control_chars(directory)
     if not os.path.isdir(directory):
         os.makedirs(directory)
 
@@ -434,7 +442,7 @@ highlight_start = "<div class=\"highlight\"><pre>"
 highlight_end = "</pre></div>"
 
 
-def process(sources, preserve_paths=True, outdir=None, language=None):
+def process(sources, preserve_paths=True, outdir=None, language=None, encoding="utf8"):
     """For each source file passed as argument, generate the documentation."""
 
     if not outdir:
@@ -447,8 +455,8 @@ def process(sources, preserve_paths=True, outdir=None, language=None):
     # Proceed to generating the documentation.
     if sources:
         outdir = ensure_directory(outdir)
-        css = open(path.join(outdir, "pycco.css"), "w")
-        css.write(pycco_styles)
+        css = open(path.join(outdir, "pycco.css"), "wb")
+        css.write(pycco_styles.encode(encoding))
         css.close()
 
         def next_file():
@@ -460,11 +468,13 @@ def process(sources, preserve_paths=True, outdir=None, language=None):
             except OSError:
                 pass
 
-            with open(dest, "w") as f:
-                f.write(generate_documentation(s, preserve_paths=preserve_paths, outdir=outdir,
-                                               language=language))
+            with open(dest, "wb") as f:
+                f.write(generate_documentation(s, preserve_paths=preserve_paths,
+                                               outdir=outdir,
+                                               language=language,
+                                               encoding=encoding))
 
-            print "pycco = {} -> {}".format(s, dest)
+            print("pycco = {} -> {}".format(s, dest))
 
             if sources:
                 next_file()
diff --git a/tests/test_pycco.py b/tests/test_pycco.py
index 04cb57e..22503c2 100644
--- a/tests/test_pycco.py
+++ b/tests/test_pycco.py
@@ -1,18 +1,24 @@
 import copy
+import os
 import tempfile
+import time
+
 import pytest
-import os
-import re
 from hypothesis import given, example, assume
 from hypothesis.strategies import lists, text, booleans, choices, none
 
 import pycco.main as p
 
+
 PYTHON = p.languages['.py']
 PYCCO_SOURCE = 'pycco/main.py'
 FOO_FUNCTION = """def foo():\n    return True"""
 
 
+def get_language(choice):
+    return choice(list(p.languages.values()))
+
+
 @given(lists(text()), text())
 def test_shift(fragments, default):
     if fragments == []:
@@ -33,7 +39,7 @@ def test_destination(filepath, preserve_paths, outdir):
 
 @given(choices(), text())
 def test_parse(choice, source):
-    l = choice(p.languages.values())
+    l = get_language(choice)
     parsed = p.parse(source, l)
     assert [{"code_text", "docs_text"} == set(s.keys()) for s in parsed]
 
@@ -69,7 +75,11 @@ def test_get_language_bad_source(source):
     with pytest.raises(ValueError) as e:
         assert p.get_language(source, "badlang")
 
-    assert e.value.message == "Can't figure out the language!"
+    msg = "Can't figure out the language!"
+    try:
+        assert e.value.message == msg
+    except AttributeError:
+        assert e.value.args[0] == msg
 
 
 @given(text() | none())
@@ -80,16 +90,13 @@ def test_get_language_bad_code(code):
 
 @given(text(max_size=64))
 def test_ensure_directory(dir_name):
-    tempdir = os.path.join(tempfile.gettempdir(), dir_name)
+    tempdir = os.path.join(tempfile.gettempdir(), str(int(time.time())), dir_name)
 
-    # Copy and paste sanitization from function, but only for housekeeping. We
+    # Use sanitization from function, but only for housekeeping. We
     # pass in the unsanitized string to the function.
-    control_chars = ''.join(map(unichr, range(0, 32) + range(127, 160)))
-    control_char_re = re.compile(u'[{}]'.format(re.escape(control_chars)))
-    safe_name = control_char_re.sub('', tempdir)
+    safe_name = p.remove_control_chars(dir_name)
 
-    if not os.path.isdir(safe_name):
-        assume(os.access(safe_name, os.W_OK))
+    if not os.path.isdir(safe_name) and os.access(safe_name, os.W_OK):
         p.ensure_directory(tempdir)
         assert os.path.isdir(safe_name)