From 8fad4a155bd2e9fde45e155b9dc35c695470ff52 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Mas=C5=82owski?= <m.maslowski@clearcode.cc>
Date: Wed, 14 Sep 2016 17:15:45 +0200
Subject: [PATCH] Add mypy type declarations, fixes #12

---
 .travis.yml                 |  5 +++++
 querystringsafe_base64.py   | 19 ++++++++++++-------
 tests/test_encode_decode.py |  7 ++++---
 3 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 96b8f10..fcffb1a 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -15,9 +15,14 @@ install:
 - "[ ! -d ~/bin ] && mkdir ~/bin || true"
 - wget -q -O - https://github.com/yihui/crandalf/raw/master/inst/scripts/install-pandoc
   | bash
+- if [[ "$TRAVIS_PYTHON_VERSION" == 3* ]]; then pip install mypy-lang; fi
 script:
 - py.test --verbose --capture=no --showlocals --cov querystringsafe_base64
 - pylama
+- if [[ "$TRAVIS_PYTHON_VERSION" == 3* ]]; then
+    mypy querystringsafe_base64.py tests;
+    mypy --py2 querystringsafe_base64.py tests;
+  fi
 after_success:
 - coveralls
 matrix:
diff --git a/querystringsafe_base64.py b/querystringsafe_base64.py
index 09bf7d3..3f24fa3 100644
--- a/querystringsafe_base64.py
+++ b/querystringsafe_base64.py
@@ -23,34 +23,38 @@
 from base64 import urlsafe_b64encode, urlsafe_b64decode
 
 __version__ = '0.2.0'
-PY2 = sys.version_info < (3, 0)
+
+if not sys.version_info < (3, 0):
+    unicode = str
 
 
 def encode(to_encode):
+    # type: (bytes) -> unicode
     """
-    Encode an arbitrary string as a base64 that is safe to put as a URL query value.
+    Encode an arbitrary bytestring as a base64 that is safe to put as a URL query value.
 
     urllib.quote and urllib.quote_plus do not have any effect on the
     result of querystringsafe_base64.encode.
 
-    :param (str, bytes) to_encode:
+    :param bytes to_encode:
     :rtype: str
     :return: a string that is safe to put as a value in an URL query
         string - like base64, except characters ['+', '/', '='] are
         replaced with ['-', '_', '.'] consequently
     """
     encoded = urlsafe_b64encode(to_encode).replace(b'=', b'.')
-    if PY2:
+    if sys.version_info < (3, 0):
         return encoded
     return encoded.decode()
 
 
 def decode(encoded):
+    # type: (unicode) -> bytes
     """
     Decode the result of querystringsafe_base64_encode or a regular base64.
 
     .. note ::
-        As a regular base64 string does not contain dots, replcing dots with
+        As a regular base64 string does not contain dots, replacing dots with
         equal signs does basically noting to it. Also,
         base64.urlsafe_b64decode allows to decode both safe and unsafe base64.
         Therefore this function may also be used to decode the regular base64.
@@ -59,6 +63,7 @@ def decode(encoded):
     :rtype: str, bytes
     :return: decoded string
     """
-    if PY2:
+    if sys.version_info < (3, 0):
         return urlsafe_b64decode(str(encoded).replace('.', '='))
-    return urlsafe_b64decode(bytes(encoded, 'ascii').replace(b'.', b'='))
+    else:
+        return urlsafe_b64decode(bytes(encoded, 'ascii').replace(b'.', b'='))
diff --git a/tests/test_encode_decode.py b/tests/test_encode_decode.py
index a095455..f11b1e4 100644
--- a/tests/test_encode_decode.py
+++ b/tests/test_encode_decode.py
@@ -21,12 +21,13 @@
 
 from base64 import b64encode, b64decode, urlsafe_b64decode
 import string
-try:
+import sys
+if sys.version_info < (3, 0):
     from urllib import quote_plus, unquote_plus
-except ImportError:
+else:
     from urllib.parse import quote_plus, unquote_plus
 
-import pytest
+import pytest  # type: ignore
 import querystringsafe_base64
 
 # We want to test querystringsafe_base64.encode with a string that normally