Basic Python 3 support

4 files changed, 35 insertions, 18 deletions

diff --git a/.travis.yml b/.travis.yml
index 8aa7660..e4997ac 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -1,7 +1,7 @@
 language: python
 python:
     - '2.7'
-    - '3.4'
+    - '3.5'
 install:
     - 'pip install -r requirements.txt'
     - 'pip install -r requirements.test.txt'
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 1c5f413..8103d01 100644
--- a/pycco/main.py
+++ b/pycco/main.py
@@ -36,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
@@ -45,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)
 
 
@@ -227,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))
@@ -404,15 +406,20 @@ def shift(list, default):
         return default
 
 
+def remove_control_chars(s):
+    from .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)
 
@@ -435,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:
@@ -448,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():
@@ -461,9 +468,11 @@ 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))
 
diff --git a/tests/test_pycco.py b/tests/test_pycco.py
index 04cb57e..90a40fd 100644
--- a/tests/test_pycco.py
+++ b/tests/test_pycco.py
@@ -12,6 +12,8 @@ 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):
@@ -33,7 +35,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 +71,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())
@@ -82,11 +88,9 @@ def test_get_language_bad_code(code):
 def test_ensure_directory(dir_name):
     tempdir = os.path.join(tempfile.gettempdir(), 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))