3939"""
4040# pylint: disable=no-value-for-parameter
4141
42+ from __future__ import annotations
43+
4244import logging
43- from typing import Collection
45+ from types import CodeType
46+ from typing import Any , Callable , Collection , TypeVar
4447
4548import jinja2
49+ from jinja2 .environment import Template
4650from wrapt import wrap_function_wrapper as _wrap
4751
4852from opentelemetry .instrumentation .instrumentor import BaseInstrumentor
4953from opentelemetry .instrumentation .jinja2 .package import _instruments
5054from opentelemetry .instrumentation .jinja2 .version import __version__
5155from opentelemetry .instrumentation .utils import unwrap
52- from opentelemetry .trace import SpanKind , get_tracer
56+ from opentelemetry .trace import SpanKind , Tracer , get_tracer
5357
5458logger = logging .getLogger (__name__ )
5559
5660ATTRIBUTE_JINJA2_TEMPLATE_NAME = "jinja2.template_name"
5761ATTRIBUTE_JINJA2_TEMPLATE_PATH = "jinja2.template_path"
5862DEFAULT_TEMPLATE_NAME = "<memory>"
5963
64+ R = TypeVar ("R" )
65+
6066
61- def _with_tracer_wrapper (func ):
67+ def _with_tracer_wrapper (
68+ func : Callable [
69+ [Tracer , Callable [..., R ], Any , list [Any ], dict [str , Any ]], R
70+ ],
71+ ) -> Callable [
72+ [Tracer ], Callable [[Callable [..., R ], Any , list [Any ], dict [str , Any ]], R ]
73+ ]:
6274 """Helper for providing tracer for wrapper functions."""
6375
64- def _with_tracer (tracer ):
65- def wrapper (wrapped , instance , args , kwargs ):
76+ def _with_tracer (
77+ tracer : Tracer ,
78+ ) -> Callable [[Callable [..., R ], Any , list [Any ], dict [str , Any ]], R ]:
79+ def wrapper (
80+ wrapped : Callable [..., R ],
81+ instance : Any ,
82+ args : list [Any ],
83+ kwargs : dict [str , Any ],
84+ ) -> R :
6685 return func (tracer , wrapped , instance , args , kwargs )
6786
6887 return wrapper
@@ -71,7 +90,13 @@ def wrapper(wrapped, instance, args, kwargs):
7190
7291
7392@_with_tracer_wrapper
74- def _wrap_render (tracer , wrapped , instance , args , kwargs ):
93+ def _wrap_render (
94+ tracer : Tracer ,
95+ wrapped : Callable [..., Any ],
96+ instance : Template ,
97+ args : list [Any ],
98+ kwargs : dict [str , Any ],
99+ ):
75100 """Wrap `Template.render()` or `Template.generate()`"""
76101 with tracer .start_as_current_span (
77102 "jinja2.render" ,
@@ -84,7 +109,13 @@ def _wrap_render(tracer, wrapped, instance, args, kwargs):
84109
85110
86111@_with_tracer_wrapper
87- def _wrap_compile (tracer , wrapped , _ , args , kwargs ):
112+ def _wrap_compile (
113+ tracer : Tracer ,
114+ wrapped : Callable [..., CodeType ],
115+ _ ,
116+ args : list [Any ],
117+ kwargs : dict [str , Any ],
118+ ) -> CodeType :
88119 with tracer .start_as_current_span (
89120 "jinja2.compile" ,
90121 kind = SpanKind .INTERNAL ,
@@ -100,7 +131,13 @@ def _wrap_compile(tracer, wrapped, _, args, kwargs):
100131
101132
102133@_with_tracer_wrapper
103- def _wrap_load_template (tracer , wrapped , _ , args , kwargs ):
134+ def _wrap_load_template (
135+ tracer : Tracer ,
136+ wrapped : Callable [..., Template ],
137+ _ ,
138+ args : list [Any ],
139+ kwargs : dict [str , Any ],
140+ ) -> Template :
104141 with tracer .start_as_current_span (
105142 "jinja2.load" ,
106143 kind = SpanKind .INTERNAL ,
@@ -128,7 +165,7 @@ class Jinja2Instrumentor(BaseInstrumentor):
128165 def instrumentation_dependencies (self ) -> Collection [str ]:
129166 return _instruments
130167
131- def _instrument (self , ** kwargs ):
168+ def _instrument (self , ** kwargs : Any ):
132169 tracer_provider = kwargs .get ("tracer_provider" )
133170 tracer = get_tracer (
134171 __name__ ,
@@ -146,7 +183,7 @@ def _instrument(self, **kwargs):
146183 _wrap_load_template (tracer ),
147184 )
148185
149- def _uninstrument (self , ** kwargs ):
186+ def _uninstrument (self , ** kwargs : Any ):
150187 unwrap (jinja2 .Template , "render" )
151188 unwrap (jinja2 .Template , "generate" )
152189 unwrap (jinja2 .Environment , "compile" )
0 commit comments