Enable mypy for ddtrace core, add type hinting for tracer and filters

ddtrace/_hooks.py

@@ -1,5 +1,7 @@
 import collections
 from copy import deepcopy
+from typing import DefaultDict
+from typing import Set
 
 from .internal.logger import get_logger
 from .vendor import attr
@@ -20,7 +22,7 @@ def on_request(span, request, response):
             pass
     """
 
-    _hooks = attr.ib(init=False, factory=lambda: collections.defaultdict(set))
+    _hooks = attr.ib(init=False, factory=lambda: collections.defaultdict(set), type=DefaultDict[str, Set])
 
     def __deepcopy__(self, memodict=None):
         hooks = Hooks()

ddtrace/compat.py

@@ -4,6 +4,9 @@
 import sys
 import textwrap
 import threading
+from typing import Any
+from typing import AnyStr
+from typing import Text
 
 from ddtrace.vendor import six
 
@@ -29,17 +32,17 @@
 PYTHON_INTERPRETER = platform.python_implementation()
 
 try:
-    StringIO = six.moves.cStringIO
+    StringIO = six.moves.cStringIO  # type: ignore[attr-defined]
 except ImportError:
     StringIO = six.StringIO
 
-httplib = six.moves.http_client
-urlencode = six.moves.urllib.parse.urlencode
-parse = six.moves.urllib.parse
-Queue = six.moves.queue.Queue
+httplib = six.moves.http_client  # type: ignore[attr-defined]
+urlencode = six.moves.urllib.parse.urlencode  # type: ignore[attr-defined]
+parse = six.moves.urllib.parse  # type: ignore[attr-defined]
+Queue = six.moves.queue.Queue  # type: ignore[attr-defined]
 iteritems = six.iteritems
 reraise = six.reraise
-reload_module = six.moves.reload_module
+reload_module = six.moves.reload_module  # type: ignore[attr-defined]
 
 stringify = six.text_type
 string_type = six.string_types[0]
@@ -52,10 +55,11 @@
 if PYTHON_VERSION_INFO >= (3, 7):
     pattern_type = re.Pattern
 else:
-    pattern_type = re._pattern_type
+    pattern_type = re._pattern_type  # type: ignore[misc,attr-defined]
 
 
 def is_integer(obj):
+    # type: (Any) -> bool
     """Helper to determine if the provided ``obj`` is an integer type or not"""
     # DEV: We have to make sure it is an integer and not a boolean
     # >>> type(True)
@@ -71,6 +75,7 @@ def is_integer(obj):
     from time import time as _time
 
     def time_ns():
+        # type: () -> int
         return int(_time() * 10e5) * 1000
 
 
@@ -85,15 +90,17 @@ def time_ns():
 except ImportError:
 
     def monotonic_ns():
+        # type: () -> int
         return int(monotonic() * 1e9)
 
 
 try:
     from time import process_time_ns
 except ImportError:
-    from time import clock as _process_time
+    from time import clock as _process_time  # type: ignore[attr-defined]
 
     def process_time_ns():
+        # type: () -> int
         return int(_process_time() * 1e9)
 
 
@@ -104,10 +111,10 @@ def process_time_ns():
 
 
 if sys.version_info.major < 3:
-    if isinstance(threading.current_thread(), threading._MainThread):
+    if isinstance(threading.current_thread(), threading._MainThread):  # type: ignore[attr-defined]
         main_thread = threading.current_thread()
     else:
-        main_thread = threading._shutdown.im_self
+        main_thread = threading._shutdown.im_self  # type: ignore[attr-defined]
 else:
     main_thread = threading.main_thread()
 
@@ -152,7 +159,7 @@ def func_wrapper(*args, **kwargs):
     # asyncio is missing so we can't have coroutines; these
     # functions are used only to ensure code executions in case
     # of an unexpected behavior
-    def iscoroutinefunction(fn):
+    def iscoroutinefunction(fn):  # type: ignore
         return False
 
     def make_async_decorator(tracer, fn, *params, **kw_params):
@@ -161,6 +168,7 @@ def make_async_decorator(tracer, fn, *params, **kw_params):
 
 # DEV: There is `six.u()` which does something similar, but doesn't have the guard around `hasattr(s, 'decode')`
 def to_unicode(s):
+    # type: (AnyStr) -> Text
     """ Return a unicode string for the given bytes or string instance. """
     # No reason to decode if we already have the unicode compatible object we expect
     # DEV: `six.text_type` will be a `str` for python 3 and `unicode` for python 2
@@ -200,7 +208,7 @@ def get_connection_response(conn):
 try:
     import contextvars  # noqa
 except ImportError:
-    from ddtrace.vendor import contextvars  # noqa
+    from ddtrace.vendor import contextvars  # type: ignore  # noqa
 
     CONTEXTVARS_IS_AVAILABLE = False
 else:

ddtrace/filters.py

@@ -9,7 +9,7 @@
 from .ext import http
 
 
-class TraceFilter(six.with_metaclass(abc.ABCMeta)):
+class TraceFilter(six.with_metaclass(abc.ABCMeta)):  # type: ignore[misc]
     @abc.abstractmethod
     def process_trace(self, trace):
         # type: (List[Span]) -> Optional[List[Span]]
@@ -52,6 +52,7 @@ def __init__(self, regexps):
         self._regexps = [re.compile(regexp) for regexp in regexps]
 
     def process_trace(self, trace):
+        # type: (List[Span]) -> Optional[List[Span]]
         """
         When the filter is registered in the tracer, process_trace is called by
         on each trace before it is sent to the agent, the returned value will

ddtrace/internal/_encoding.pyi

@@ -0,0 +1,6 @@
+from typing import Any
+from typing import Union
+
+class MsgpackEncoder(object):
+    content_type: str
+    def _decode(self, data: Union[str, bytes]) -> Any: ...

ddtrace/internal/_rand.pyi

@@ -0,0 +1,2 @@
+def seed() -> None: ...
+def rand64bits(check_pid: bool = True) -> int: ...

ddtrace/internal/writer.py

@@ -133,7 +133,7 @@ def stop(self, timeout=None):
 
     @abc.abstractmethod
     def write(self, trace):
-        # type: (List[Span]) -> None
+        # type: (Optional[List[Span]]) -> None
         pass
 
 
@@ -159,7 +159,7 @@ def stop(self, timeout=None):
         return
 
     def write(self, spans):
-        # type: (List[Span]) -> None
+        # type: (Optional[List[Span]]) -> None
         if not spans:
             return
 

ddtrace/provider.py

@@ -9,7 +9,7 @@
 _DD_CONTEXTVAR = contextvars.ContextVar("datadog_contextvar", default=None)
 
 
-class BaseContextProvider(six.with_metaclass(abc.ABCMeta)):
+class BaseContextProvider(six.with_metaclass(abc.ABCMeta)):  # type: ignore[misc]
     """
     A ``ContextProvider`` is an interface that provides the blueprint
     for a callable class, capable to retrieve the current active

ddtrace/sampler.py

@@ -26,13 +26,13 @@
 KNUTH_FACTOR = 1111111111111111111
 
 
-class BaseSampler(six.with_metaclass(abc.ABCMeta)):
+class BaseSampler(six.with_metaclass(abc.ABCMeta)):  # type: ignore[misc]
     @abc.abstractmethod
     def sample(self, span):
         pass
 
 
-class BasePrioritySampler(six.with_metaclass(abc.ABCMeta)):
+class BasePrioritySampler(six.with_metaclass(abc.ABCMeta)):  # type: ignore[misc]
     @abc.abstractmethod
     def update_rate_by_service_sample_rates(self, sample_rates):
         pass