Add type hints and expose them to users

Author: qnighy

MANIFEST.in

@@ -8,3 +8,4 @@ include .coveragerc
 include .editorconfig
 include Makefile
 include requirements.txt
+include src/dotenv/py.typed

requirements.txt

@@ -1,4 +1,5 @@
 bumpversion
+typing
 click
 flake8>=2.2.3
 ipython

setup.py

@@ -29,6 +29,12 @@
               'configurations', 'python'],
     packages=['dotenv'],
     package_dir={'': 'src'},
+    package_data={
+        'dotenv': ['py.typed'],
+    },
+    install_requires=[
+        'typing',
+    ],
     extras_require={
         'cli': ['click>=5.0', ],
     },

src/dotenv/__init__.py

@@ -1,12 +1,15 @@
+from typing import Any, Optional
 from .main import load_dotenv, get_key, set_key, unset_key, find_dotenv, dotenv_values
 
 
 def load_ipython_extension(ipython):
+    # type: (Any) -> None
     from .ipython import load_ipython_extension
     load_ipython_extension(ipython)
 
 
 def get_cli_string(path=None, action=None, key=None, value=None, quote=None):
+    # type: (Optional[str], Optional[str], Optional[str], Optional[str], Optional[str]) -> str
     """Returns a string suitable for running as a shell script.
 
     Useful for converting a arguments passed to a fabric task

src/dotenv/cli.py

@@ -1,5 +1,6 @@
 import os
 import sys
+from typing import Any, List
 
 try:
     import click
@@ -22,6 +23,7 @@
 @click.version_option(version=__version__)
 @click.pass_context
 def cli(ctx, file, quote):
+    # type: (click.Context, Any, Any) -> None
     '''This script is used to set, get or unset values from a .env file.'''
     ctx.obj = {}
     ctx.obj['FILE'] = file
@@ -31,6 +33,7 @@ def cli(ctx, file, quote):
 @cli.command()
 @click.pass_context
 def list(ctx):
+    # type: (click.Context) -> None
     '''Display all the stored key/value.'''
     file = ctx.obj['FILE']
     dotenv_as_dict = dotenv_values(file)
@@ -43,6 +46,7 @@ def list(ctx):
 @click.argument('key', required=True)
 @click.argument('value', required=True)
 def set(ctx, key, value):
+    # type: (click.Context, Any, Any) -> None
     '''Store the given key/value.'''
     file = ctx.obj['FILE']
     quote = ctx.obj['QUOTE']
@@ -57,6 +61,7 @@ def set(ctx, key, value):
 @click.pass_context
 @click.argument('key', required=True)
 def get(ctx, key):
+    # type: (click.Context, Any) -> None
     '''Retrieve the value for the given key.'''
     file = ctx.obj['FILE']
     stored_value = get_key(file, key)
@@ -70,6 +75,7 @@ def get(ctx, key):
 @click.pass_context
 @click.argument('key', required=True)
 def unset(ctx, key):
+    # type: (click.Context, Any) -> None
     '''Removes the given key.'''
     file = ctx.obj['FILE']
     quote = ctx.obj['QUOTE']
@@ -84,13 +90,14 @@ def unset(ctx, key):
 @click.pass_context
 @click.argument('commandline', nargs=-1, type=click.UNPROCESSED)
 def run(ctx, commandline):
+    # type: (click.Context, List[str]) -> None
     """Run command with environment variables present."""
     file = ctx.obj['FILE']
     dotenv_as_dict = dotenv_values(file)
     if not commandline:
         click.echo('No command given.')
         exit(1)
-    ret = run_command(commandline, dotenv_as_dict)
+    ret = run_command(commandline, dotenv_as_dict)  # type: ignore
     exit(ret)
 
 

src/dotenv/compat.py

@@ -1,9 +1,10 @@
+from typing import Text, Type
 import sys
-try:
-    from StringIO import StringIO  # noqa
-except ImportError:
+if sys.version_info >= (3, 0):
     from io import StringIO  # noqa
+else:
+    from StringIO import StringIO  # noqa
 
-PY2 = sys.version_info[0] == 2
-WIN = sys.platform.startswith('win')
-text_type = unicode if PY2 else str  # noqa
+PY2 = sys.version_info[0] == 2  # type: bool
+WIN = sys.platform.startswith('win')  # type: bool
+text_type = Text  # type: Type[Text]

src/dotenv/ipython.py

@@ -1,8 +1,8 @@
 from __future__ import print_function
 
-from IPython.core.magic import Magics, line_magic, magics_class
-from IPython.core.magic_arguments import (argument, magic_arguments,
-                                          parse_argstring)
+from IPython.core.magic import Magics, line_magic, magics_class  # type: ignore
+from IPython.core.magic_arguments import (argument, magic_arguments,  # type: ignore
+                                          parse_argstring)  # type: ignore
 
 from .main import find_dotenv, load_dotenv
 

src/dotenv/main.py

@@ -9,13 +9,26 @@
 import sys
 from subprocess import Popen
 import tempfile
+from typing import (Any, Dict, Iterator, List, Match, NamedTuple, Optional,  # noqa
+                    Pattern, Union, TYPE_CHECKING, Text, IO, Tuple)  # noqa
 import warnings
-from collections import OrderedDict, namedtuple
+from collections import OrderedDict
 from contextlib import contextmanager
 
 from .compat import StringIO, PY2, WIN, text_type
 
-__posix_variable = re.compile(r'\$\{[^\}]*\}')
+if TYPE_CHECKING:  # pragma: no cover
+    if sys.version_info >= (3, 6):
+        _PathLike = os.PathLike
+    else:
+        _PathLike = Text
+
+    if sys.version_info >= (3, 0):
+        _StringIO = StringIO
+    else:
+        _StringIO = StringIO[Text]
+
+__posix_variable = re.compile(r'\$\{[^\}]*\}')  # type: Pattern[Text]
 
 _binding = re.compile(
     r"""
@@ -42,30 +55,37 @@
         )
     """.format(r'[^\S\r\n]'),
     re.MULTILINE | re.VERBOSE,
-)
+)  # type: Pattern[Text]
 
-_escape_sequence = re.compile(r"\\[\\'\"abfnrtv]")
+_escape_sequence = re.compile(r"\\[\\'\"abfnrtv]")  # type: Pattern[Text]
 
 
-Binding = namedtuple('Binding', 'key value original')
+Binding = NamedTuple("Binding", [("key", Optional[Text]),
+                                 ("value", Optional[Text]),
+                                 ("original", Text)])
 
 
 def decode_escapes(string):
+    # type: (Text) -> Text
     def decode_match(match):
-        return codecs.decode(match.group(0), 'unicode-escape')
+        # type: (Match[Text]) -> Text
+        return codecs.decode(match.group(0), 'unicode-escape')  # type: ignore
 
     return _escape_sequence.sub(decode_match, string)
 
 
 def is_surrounded_by(string, char):
+    # type: (Text, Text) -> bool
     return (
         len(string) > 1
         and string[0] == string[-1] == char
     )
 
 
 def parse_binding(string, position):
+    # type: (Text, int) -> Tuple[Binding, int]
     match = _binding.match(string, position)
+    assert match is not None
     (matched, key, value) = match.groups()
     if key is None or value is None:
         key = None
@@ -80,6 +100,7 @@ def parse_binding(string, position):
 
 
 def parse_stream(stream):
+    # type:(IO[Text]) -> Iterator[Binding]
     string = stream.read()
     position = 0
     length = len(string)
@@ -91,23 +112,26 @@ def parse_stream(stream):
 class DotEnv():
 
     def __init__(self, dotenv_path, verbose=False):
-        self.dotenv_path = dotenv_path
-        self._dict = None
-        self.verbose = verbose
+        # type: (Union[Text, _PathLike, _StringIO], bool) -> None
+        self.dotenv_path = dotenv_path  # type: Union[Text,_PathLike, _StringIO]
+        self._dict = None  # type: Optional[Dict[Text, Text]]
+        self.verbose = verbose  # type: bool
 
     @contextmanager
     def _get_stream(self):
+        # type: () -> Iterator[IO[Text]]
         if isinstance(self.dotenv_path, StringIO):
             yield self.dotenv_path
         elif os.path.isfile(self.dotenv_path):
             with io.open(self.dotenv_path) as stream:
                 yield stream
         else:
             if self.verbose:
-                warnings.warn("File doesn't exist {}".format(self.dotenv_path))
+                warnings.warn("File doesn't exist {}".format(self.dotenv_path))  # type: ignore
             yield StringIO('')
 
     def dict(self):
+        # type: () -> Dict[Text, Text]
         """Return dotenv as dict"""
         if self._dict:
             return self._dict
@@ -117,12 +141,14 @@ def dict(self):
         return self._dict
 
     def parse(self):
+        # type: () -> Iterator[Tuple[Text, Text]]
         with self._get_stream() as stream:
             for mapping in parse_stream(stream):
                 if mapping.key is not None and mapping.value is not None:
                     yield mapping.key, mapping.value
 
     def set_as_environment_variables(self, override=False):
+        # type: (bool) -> bool
         """
         Load the current dotenv as system environemt variable.
         """
@@ -135,11 +161,12 @@ def set_as_environment_variables(self, override=False):
                 if isinstance(k, text_type) or isinstance(v, text_type):
                     k = k.encode('ascii')
                     v = v.encode('ascii')
-            os.environ[k] = v
+            os.environ[k] = v  # type: ignore
 
         return True
 
     def get(self, key):
+        # type: (Text) -> Optional[Text]
         """
         """
         data = self.dict()
@@ -148,10 +175,13 @@ def get(self, key):
             return data[key]
 
         if self.verbose:
-            warnings.warn("key %s not found in %s." % (key, self.dotenv_path))
+            warnings.warn("key %s not found in %s." % (key, self.dotenv_path))  # type: ignore
+
+        return None
 
 
 def get_key(dotenv_path, key_to_get):
+    # type: (Union[Text, _PathLike], Text) -> Optional[Text]
     """
     Gets the value of a given key from the given .env
 
@@ -162,10 +192,11 @@ def get_key(dotenv_path, key_to_get):
 
 @contextmanager
 def rewrite(path):
+    # type: (_PathLike) -> Iterator[Tuple[IO[Text], IO[Text]]]
     try:
         with tempfile.NamedTemporaryFile(mode="w+", delete=False) as dest:
             with io.open(path) as source:
-                yield (source, dest)
+                yield (source, dest)  # type: ignore
     except BaseException:
         if os.path.isfile(dest.name):
             os.unlink(dest.name)
@@ -175,6 +206,7 @@ def rewrite(path):
 
 
 def set_key(dotenv_path, key_to_set, value_to_set, quote_mode="always"):
+    # type: (_PathLike, Text, Text, Text) -> Tuple[Optional[bool], Text, Text]
     """
     Adds or Updates a key/value to the given .env
 
@@ -183,7 +215,7 @@ def set_key(dotenv_path, key_to_set, value_to_set, quote_mode="always"):
     """
     value_to_set = value_to_set.strip("'").strip('"')
     if not os.path.exists(dotenv_path):
-        warnings.warn("can't write to %s - it doesn't exist." % dotenv_path)
+        warnings.warn("can't write to %s - it doesn't exist." % dotenv_path)  # type: ignore
         return None, key_to_set, value_to_set
 
     if " " in value_to_set:
@@ -207,14 +239,15 @@ def set_key(dotenv_path, key_to_set, value_to_set, quote_mode="always"):
 
 
 def unset_key(dotenv_path, key_to_unset, quote_mode="always"):
+    # type: (_PathLike, Text, Text) -> Tuple[Optional[bool], Text]
     """
     Removes a given key from the given .env
 
     If the .env path given doesn't exist, fails
     If the given key doesn't exist in the .env, fails
     """
     if not os.path.exists(dotenv_path):
-        warnings.warn("can't delete from %s - it doesn't exist." % dotenv_path)
+        warnings.warn("can't delete from %s - it doesn't exist." % dotenv_path)  # type: ignore
         return None, key_to_unset
 
     removed = False
@@ -226,14 +259,16 @@ def unset_key(dotenv_path, key_to_unset, quote_mode="always"):
                 dest.write(mapping.original)
 
     if not removed:
-        warnings.warn("key %s not removed from %s - key doesn't exist." % (key_to_unset, dotenv_path))
+        warnings.warn("key %s not removed from %s - key doesn't exist." % (key_to_unset, dotenv_path))  # type: ignore
         return None, key_to_unset
 
     return removed, key_to_unset
 
 
 def resolve_nested_variables(values):
+    # type: (Dict[Text, Text]) -> Dict[Text, Text]
     def _replacement(name):
+        # type: (Text) -> Text
         """
         get appropriate value for a variable name.
         first search in environ, if not found,
@@ -243,6 +278,7 @@ def _replacement(name):
         return ret
 
     def _re_sub_callback(match_object):
+        # type: (Match[Text]) -> Text
         """
         From a match object gets the variable name and returns
         the correct replacement
@@ -258,6 +294,7 @@ def _re_sub_callback(match_object):
 
 
 def _walk_to_root(path):
+    # type: (Text) -> Iterator[Text]
     """
     Yield directories starting from the given directory up to the root
     """
@@ -276,6 +313,7 @@ def _walk_to_root(path):
 
 
 def find_dotenv(filename='.env', raise_error_if_not_found=False, usecwd=False):
+    # type: (Text, bool, bool) -> Text
     """
     Search in increasingly higher folders for the given file
 
@@ -312,16 +350,19 @@ def find_dotenv(filename='.env', raise_error_if_not_found=False, usecwd=False):
 
 
 def load_dotenv(dotenv_path=None, stream=None, verbose=False, override=False):
+    # type: (Union[Text, _PathLike, None], Optional[_StringIO], bool, bool) -> bool
     f = dotenv_path or stream or find_dotenv()
     return DotEnv(f, verbose=verbose).set_as_environment_variables(override=override)
 
 
 def dotenv_values(dotenv_path=None, stream=None, verbose=False):
+    # type: (Union[Text, _PathLike, None], Optional[_StringIO], bool) -> Dict[Text, Text]
     f = dotenv_path or stream or find_dotenv()
     return DotEnv(f, verbose=verbose).dict()
 
 
 def run_command(command, env):
+    # type: (List[str], Dict[str, str]) -> int
     """Run command in sub process.
 
     Runs the command in a sub process with the variables from `env`

src/dotenv/py.typed

@@ -0,0 +1 @@
+# Marker file for PEP 561