diff --git a/ddtrace/_hooks.py b/ddtrace/_hooks.py index 5ef50cc0f84..51ab58e2743 100644 --- a/ddtrace/_hooks.py +++ b/ddtrace/_hooks.py @@ -1,6 +1,9 @@ import collections from copy import deepcopy +from typing import Any +from typing import Callable from typing import DefaultDict +from typing import Optional from typing import Set from .internal.logger import get_logger @@ -29,7 +32,12 @@ def __deepcopy__(self, memodict=None): hooks._hooks = deepcopy(self._hooks, memodict) return hooks - def register(self, hook, func=None): + def register( + self, + hook, # type: Any + func=None, # type: Optional[Callable] + ): + # type: (...) -> Optional[Callable[..., Any]] """ Function used to register a hook for the provided name. @@ -63,6 +71,7 @@ def wrapper(func): return wrapper self._hooks[hook].add(func) + return None # Provide shorthand `on` method for `register` # >>> @config.falcon.hooks.on('request') @@ -70,7 +79,12 @@ def wrapper(func): # pass on = register - def deregister(self, hook, func): + def deregister( + self, + hook, # type: Any + func, # type: Callable + ): + # type: (...) -> None """ Function to deregister a function from a hook it was registered under @@ -93,7 +107,13 @@ def on_request(span, request, response): except KeyError: pass - def emit(self, hook, *args, **kwargs): + def emit( + self, + hook, # type: Any + *args, # type: Any + **kwargs # type: Any + ): + # type: (...) -> None """ Function used to call registered hook functions. diff --git a/ddtrace/_worker.py b/ddtrace/_worker.py index b9ed478501a..ea9a106f1b5 100644 --- a/ddtrace/_worker.py +++ b/ddtrace/_worker.py @@ -1,4 +1,5 @@ import threading +from typing import Optional from .internal.logger import get_logger @@ -19,7 +20,13 @@ class PeriodicWorkerThread(object): _DEFAULT_INTERVAL = 1.0 - def __init__(self, interval=_DEFAULT_INTERVAL, name=None, daemon=True): + def __init__( + self, + interval=_DEFAULT_INTERVAL, # type: float + name=None, # type: Optional[str] + daemon=True, # type: bool + ): + # type: (...) -> None """Create a new worker thread that runs a function periodically. :param interval: The interval in seconds to wait between calls to `run_periodic`. @@ -34,37 +41,45 @@ def __init__(self, interval=_DEFAULT_INTERVAL, name=None, daemon=True): self.interval = interval def start(self): + # type: () -> None """Start the periodic worker.""" _LOG.debug("Starting %s thread", self._thread.name) self._thread.start() self.started = True def stop(self): + # type: () -> None """Stop the worker.""" _LOG.debug("Stopping %s thread", self._thread.name) self._stop.set() def is_alive(self): + # type: () -> bool return self._thread.is_alive() def join(self, timeout=None): + # type: (Optional[float]) -> None return self._thread.join(timeout) def _target(self): + # type: () -> None while not self._stop.wait(self.interval): self.run_periodic() self._on_shutdown() @staticmethod def run_periodic(): + # type: () -> None """Method executed every interval.""" pass def _on_shutdown(self): + # type: () -> None _LOG.debug("Shutting down %s thread", self._thread.name) self.on_shutdown() @staticmethod def on_shutdown(): + # type: () -> None """Method ran on worker shutdown.""" pass diff --git a/ddtrace/encoding.py b/ddtrace/encoding.py index b468dfa6957..ee91b4dafb1 100644 --- a/ddtrace/encoding.py +++ b/ddtrace/encoding.py @@ -1,9 +1,18 @@ import json +from typing import Any +from typing import Dict +from typing import List +from typing import Optional +from typing import TYPE_CHECKING from .internal._encoding import MsgpackEncoder from .internal.logger import get_logger +if TYPE_CHECKING: + from .span import Span + + log = get_logger(__name__) @@ -13,6 +22,7 @@ class _EncoderBase(object): """ def encode_traces(self, traces): + # type: (List[List[Span]]) -> str """ Encodes a list of traces, expecting a list of items where each items is a list of spans. Before dump the string in a serialized format all @@ -25,6 +35,7 @@ def encode_traces(self, traces): return self.encode(normalized_traces) def encode_trace(self, trace): + # type: (List[Span]) -> str """ Encodes a trace, expecting a list of spans. Before dump the string in a serialized format all traces are normalized, calling the ``to_dict()`` method. @@ -53,10 +64,12 @@ class JSONEncoder(_EncoderBase): @staticmethod def encode(obj): + # type: (Any) -> str return json.dumps(obj) @staticmethod def join_encoded(objs): + # type: (List[str]) -> str """Join a list of encoded objects together as a json array""" return "[" + ",".join(objs) + "]" @@ -69,19 +82,23 @@ class JSONEncoderV2(JSONEncoder): content_type = "application/json" def encode_traces(self, traces): + # type: (List[List[Span]]) -> str normalized_traces = [[JSONEncoderV2._convert_span(span) for span in trace] for trace in traces] return self.encode({"traces": normalized_traces}) def encode_trace(self, trace): + # type: (List[Span]) -> str return self.encode([JSONEncoderV2._convert_span(span) for span in trace]) @staticmethod def join_encoded(objs): + # type: (List[str]) -> str """Join a list of encoded objects together as a json array""" return '{"traces":[' + ",".join(objs) + "]}" @staticmethod def _convert_span(span): + # type: (Span) -> Dict[str, Any] sp = span.to_dict() sp["trace_id"] = JSONEncoderV2._encode_id_to_hex(sp.get("trace_id")) sp["parent_id"] = JSONEncoderV2._encode_id_to_hex(sp.get("parent_id")) @@ -90,12 +107,14 @@ def _convert_span(span): @staticmethod def _encode_id_to_hex(dd_id): + # type: (Optional[int]) -> str if not dd_id: return "0000000000000000" return "%0.16X" % int(dd_id) @staticmethod def _decode_id_to_hex(hex_id): + # type: (Optional[str]) -> int if not hex_id: return 0 return int(hex_id, 16) diff --git a/ddtrace/internal/_encoding.pyi b/ddtrace/internal/_encoding.pyi index a49baba15c1..dbab4a89d0a 100644 --- a/ddtrace/internal/_encoding.pyi +++ b/ddtrace/internal/_encoding.pyi @@ -1,6 +1,9 @@ from typing import Any +from typing import List from typing import Union class MsgpackEncoder(object): content_type: str def _decode(self, data: Union[str, bytes]) -> Any: ... + def encode_trace(self, trace: List[Any]) -> bytes: ... + def join_encoded(self, objs: List[bytes]) -> bytes: ... diff --git a/ddtrace/payload.py b/ddtrace/payload.py index 64867bdafe4..4f7bd40df15 100644 --- a/ddtrace/payload.py +++ b/ddtrace/payload.py @@ -1,6 +1,15 @@ +from typing import List +from typing import Optional +from typing import TYPE_CHECKING +from typing import Union + from .encoding import Encoder +if TYPE_CHECKING: + from .span import Span + + class PayloadFull(Exception): """The payload is full.""" @@ -24,7 +33,12 @@ class Payload(object): # 5 MB should be a good average efficient size DEFAULT_MAX_PAYLOAD_SIZE = 5 * 1000000 - def __init__(self, encoder=None, max_payload_size=DEFAULT_MAX_PAYLOAD_SIZE): + def __init__( + self, + encoder=None, # type: Optional[Encoder] + max_payload_size=DEFAULT_MAX_PAYLOAD_SIZE, # type: int + ): + # type: (...) -> None """ Constructor for Payload @@ -35,10 +49,11 @@ def __init__(self, encoder=None, max_payload_size=DEFAULT_MAX_PAYLOAD_SIZE): """ self.max_payload_size = max_payload_size self.encoder = encoder or Encoder() - self.traces = [] + self.traces = [] # type: List[bytes] self.size = 0 def add_trace(self, trace): + # type: (Optional[List[Span]]) -> None """ Encode and append a trace to this payload @@ -58,6 +73,7 @@ def add_trace(self, trace): @property def length(self): + # type: () -> int """ Get the number of traces in this payload @@ -68,6 +84,7 @@ def length(self): @property def empty(self): + # type: () -> bool """ Whether this payload is empty or not @@ -77,6 +94,7 @@ def empty(self): return self.length == 0 def get_payload(self): + # type: () -> Union[str, bytes] """ Get the fully encoded payload @@ -87,6 +105,7 @@ def get_payload(self): return self.encoder.join_encoded(self.traces) def __repr__(self): + # type: () -> str """Get the string representation of this payload""" return "{0}(length={1}, size={2} B, max_payload_size={3} B)".format( self.__class__.__name__, self.length, self.size, self.max_payload_size diff --git a/ddtrace/pin.py b/ddtrace/pin.py index 77d0f782d9a..a12245ef206 100644 --- a/ddtrace/pin.py +++ b/ddtrace/pin.py @@ -1,3 +1,8 @@ +from typing import Any +from typing import Dict +from typing import Optional +from typing import TYPE_CHECKING + import ddtrace from ddtrace.vendor import debtcollector @@ -5,6 +10,10 @@ from .vendor import wrapt +if TYPE_CHECKING: + from .tracer import Tracer + + log = get_logger(__name__) @@ -29,21 +38,31 @@ class Pin(object): __slots__ = ["app", "tags", "tracer", "_target", "_config", "_initialized"] @debtcollector.removals.removed_kwarg("app_type") - def __init__(self, service=None, app=None, app_type=None, tags=None, tracer=None, _config=None): + def __init__( + self, + service=None, # type: Optional[str] + app=None, # type: Optional[str] + app_type=None, + tags=None, # type: Optional[Dict[str, str]] + tracer=None, # type: Optional[Tracer] + _config=None, # type: Optional[Dict[str, Any]] + ): + # type: (...) -> None tracer = tracer or ddtrace.tracer self.app = app self.tags = tags self.tracer = tracer - self._target = None + self._target = None # type: Optional[int] # keep the configuration attribute internal because the # public API to access it is not the Pin class - self._config = _config or {} + self._config = _config or {} # type: Dict[str, Any] # [Backward compatibility]: service argument updates the `Pin` config self._config["service_name"] = service self._initialized = True @property def service(self): + # type: () -> str """Backward compatibility: accessing to `pin.service` returns the underlying configuration value. """ @@ -59,6 +78,7 @@ def __repr__(self): @staticmethod def _find(*objs): + # type: (Any) -> Optional[Pin] """ Return the first :class:`ddtrace.pin.Pin` found on any of the provided objects or `None` if none were found @@ -78,6 +98,7 @@ def _find(*objs): @staticmethod def get_from(obj): + # type: (Any) -> Pin """Return the pin associated with the given object. If a pin is attached to `obj` but the instance is not the owner of the pin, a new pin is cloned and attached. This ensures that a pin inherited from a class is a copy for the new @@ -103,7 +124,16 @@ def get_from(obj): @classmethod @debtcollector.removals.removed_kwarg("app_type") - def override(cls, obj, service=None, app=None, app_type=None, tags=None, tracer=None): + def override( + cls, + obj, # type: Any + service=None, # type: Optional[str] + app=None, # type: Optional[str] + app_type=None, + tags=None, # type: Optional[Dict[str, str]] + tracer=None, # type: Optional[Tracer] + ): + # type: (...) -> None """Override an object with the given attributes. That's the recommended way to customize an already instrumented client, without @@ -128,10 +158,12 @@ def override(cls, obj, service=None, app=None, app_type=None, tags=None, tracer= ).onto(obj) def enabled(self): + # type: () -> bool """Return true if this pin's tracer is enabled. """ return bool(self.tracer) and self.tracer.enabled def onto(self, obj, send=True): + # type: (Any, bool) -> None """Patch this pin onto the given object. If send is true, it will also queue the metadata to be sent to the server. """ @@ -149,6 +181,7 @@ def onto(self, obj, send=True): log.debug("can't pin onto object. skipping", exc_info=True) def remove_from(self, obj): + # type: (Any) -> None # Remove pin from the object. try: pin_name = _DD_PIN_PROXY_NAME if isinstance(obj, wrapt.ObjectProxy) else _DD_PIN_NAME @@ -160,7 +193,15 @@ def remove_from(self, obj): log.debug("can't remove pin from object. skipping", exc_info=True) @debtcollector.removals.removed_kwarg("app_type") - def clone(self, service=None, app=None, app_type=None, tags=None, tracer=None): + def clone( + self, + service=None, # type: Optional[str] + app=None, # type: Optional[str] + app_type=None, + tags=None, # type: Optional[Dict[str, str]] + tracer=None, # type: Optional[Tracer] + ): + # type: (...) -> Pin """Return a clone of the pin with the given attributes replaced.""" # do a shallow copy of Pin dicts if not tags and self.tags: diff --git a/ddtrace/sampler.py b/ddtrace/sampler.py index 1eed2115168..52762f6087e 100644 --- a/ddtrace/sampler.py +++ b/ddtrace/sampler.py @@ -3,6 +3,11 @@ Any `sampled = False` trace won't be written, and can be ignored by the instrumentation. """ import abc +from typing import Any +from typing import Dict +from typing import List +from typing import Optional +from typing import TYPE_CHECKING from .compat import iteritems from .compat import pattern_type @@ -18,6 +23,10 @@ from .vendor import six +if TYPE_CHECKING: + from .span import Span + + log = get_logger(__name__) MAX_TRACE_ID = 2 ** 64 @@ -42,6 +51,7 @@ class AllSampler(BaseSampler): """Sampler sampling all the traces""" def sample(self, span): + # type: (Span) -> bool return True @@ -53,6 +63,7 @@ class RateSampler(BaseSampler): """ def __init__(self, sample_rate=1): + # type: (float) -> None if sample_rate < 0: raise ValueError("sample_rate of {} is negative".format(sample_rate)) elif sample_rate > 1: @@ -63,10 +74,12 @@ def __init__(self, sample_rate=1): log.debug("initialized RateSampler, sample %s%% of traces", 100 * sample_rate) def set_sample_rate(self, sample_rate): + # type: (float) -> None self.sample_rate = float(sample_rate) self.sampling_id_threshold = self.sample_rate * MAX_TRACE_ID def sample(self, span): + # type: (Span) -> bool return ((span.trace_id * KNUTH_FACTOR) % MAX_TRACE_ID) <= self.sampling_id_threshold @@ -78,24 +91,37 @@ class RateByServiceSampler(BaseSampler, BasePrioritySampler): """ @staticmethod - def _key(service=None, env=None): + def _key( + service=None, # type: Optional[str] + env=None, # type: Optional[str] + ): + # type: (...) -> str """Compute a key with the same format used by the Datadog agent API.""" service = service or "" env = env or "" return "service:" + service + ",env:" + env def __init__(self, sample_rate=1): + # type: (float) -> None self.sample_rate = sample_rate self._by_service_samplers = self._get_new_by_service_sampler() def _get_new_by_service_sampler(self): + # type: () -> Dict[str, RateSampler] return {self._default_key: RateSampler(self.sample_rate)} - def set_sample_rate(self, sample_rate, service="", env=""): + def set_sample_rate( + self, + sample_rate, # type: float + service="", # type: str + env="", # type: str + ): + # type: (...) -> None self._by_service_samplers[self._key(service, env)] = RateSampler(sample_rate) def sample(self, span): - tags = span.tracer.tags + # type: (Span) -> bool + tags = span.tracer.tags # type: ignore[union-attr] env = tags[ENV_KEY] if ENV_KEY in tags else None key = self._key(span.service, env) @@ -104,6 +130,7 @@ def sample(self, span): return sampler.sample(span) def update_rate_by_service_sample_rates(self, rate_by_service): + # type: (Dict[str, float]) -> None new_by_service_samplers = self._get_new_by_service_sampler() for key, sample_rate in iteritems(rate_by_service): new_by_service_samplers[key] = RateSampler(sample_rate) @@ -121,7 +148,13 @@ class DatadogSampler(BaseSampler, BasePrioritySampler): NO_RATE_LIMIT = -1 DEFAULT_RATE_LIMIT = 100 - def __init__(self, rules=None, default_sample_rate=None, rate_limit=None): + def __init__( + self, + rules=None, # type: Optional[List[SamplingRule]] + default_sample_rate=None, # type: Optional[float] + rate_limit=None, # type: Optional[int] + ): + # type: (...) -> None """ Constructor for DatadogSampler sampler @@ -162,16 +195,19 @@ def __init__(self, rules=None, default_sample_rate=None, rate_limit=None): self.default_sampler = SamplingRule(sample_rate=default_sample_rate) def update_rate_by_service_sample_rates(self, sample_rates): + # type: (Dict[str, float]) -> None # Pass through the call to our RateByServiceSampler if isinstance(self.default_sampler, RateByServiceSampler): self.default_sampler.update_rate_by_service_sample_rates(sample_rates) def _set_priority(self, span, priority): + # type: (Span, int) -> None if span._context: span._context.sampling_priority = priority span.sampled = priority is AUTO_KEEP def sample(self, span): + # type: (Span) -> bool """ Decide whether the provided span should be sampled or not @@ -236,7 +272,13 @@ class SamplingRule(BaseSampler): NO_RULE = object() - def __init__(self, sample_rate, service=NO_RULE, name=NO_RULE): + def __init__( + self, + sample_rate, # type: float + service=NO_RULE, # type: Any + name=NO_RULE, # type: Any + ): + # type: (...) -> None """ Configure a new :class:`SamplingRule` @@ -275,10 +317,12 @@ def __init__(self, sample_rate, service=NO_RULE, name=NO_RULE): @property def sample_rate(self): + # type: () -> float return self._sample_rate @sample_rate.setter def sample_rate(self, sample_rate): + # type: (float) -> None self._sample_rate = sample_rate self._sampling_id_threshold = sample_rate * MAX_TRACE_ID @@ -312,6 +356,7 @@ def _pattern_matches(self, prop, pattern): return prop == pattern def matches(self, span): + # type: (Span) -> bool """ Return if this span matches this rule @@ -329,6 +374,7 @@ def matches(self, span): ) def sample(self, span): + # type: (Span) -> bool """ Return if this rule chooses to sample the span