fix typing

Author: mabdinur

ddtrace/_trace/_span_link.py

@@ -13,7 +13,7 @@
 Usage
 -----
 
-SpanLinks can be set using :meth:`ddtrace.Span.link_span(...)` Ex::
+SpanLinks can be set using :meth:`ddtrace._trace.span.Span.link_span(...)` Ex::
 
     from ddtrace.trace import tracer
 

ddtrace/appsec/_capabilities.py

@@ -59,7 +59,7 @@ def _asm_feature_is_required() -> bool:
     return (_FEATURE_REQUIRED & flags) != 0
 
 
-def _rc_capabilities(test_tracer: Optional[ddtrace.Tracer] = None) -> Flags:
+def _rc_capabilities(test_tracer: Optional[ddtrace._trace.tracer.Tracer] = None) -> Flags:
     tracer = ddtrace.tracer if test_tracer is None else test_tracer
     value = Flags(0)
     if ddtrace.config._remote_config_enabled:
@@ -74,7 +74,7 @@ def _rc_capabilities(test_tracer: Optional[ddtrace.Tracer] = None) -> Flags:
     return value
 
 
-def _appsec_rc_capabilities(test_tracer: Optional[ddtrace.Tracer] = None) -> str:
+def _appsec_rc_capabilities(test_tracer: Optional[ddtrace._trace.tracer.Tracer] = None) -> str:
     r"""return the bit representation of the composed capabilities in base64
     bit 0: Reserved
     bit 1: ASM 1-click Activation

ddtrace/contrib/internal/logging/patch.py

@@ -80,7 +80,7 @@ def _w_makeRecord(func, instance, args, kwargs):
     # logs from internal logger may explicitly pass the current span to
     # avoid deadlocks in getting the current span while already in locked code.
     span_from_log = getattr(record, _LOG_SPAN_KEY, None)
-    if isinstance(span_from_log, ddtrace.Span):
+    if isinstance(span_from_log, ddtrace._trace.span.Span):
         span = span_from_log
 
     if tracer:

ddtrace/contrib/internal/sqlalchemy/engine.py

@@ -37,7 +37,7 @@ def trace_engine(engine, tracer=None, service=None):
     Add tracing instrumentation to the given sqlalchemy engine or instance.
 
     :param sqlalchemy.Engine engine: a SQLAlchemy engine class or instance
-    :param ddtrace.Tracer tracer: a tracer instance. will default to the global
+    :param ddtrace._trace.tracer.Tracer tracer: a tracer instance. will default to the global
     :param str service: the name of the service to trace.
     """
     tracer = tracer or ddtrace.tracer  # by default use global

ddtrace/contrib/internal/trace_utils.py

@@ -257,7 +257,7 @@ def _store_request_headers(headers, span, integration_config):
     :param headers: All the request's http headers, will be filtered through the whitelist
     :type headers: dict or list
     :param span: The Span instance where tags will be stored
-    :type span: ddtrace.Span
+    :type span: ddtrace._trace.span.Span
     :param integration_config: An integration specific config object.
     :type integration_config: ddtrace.settings.IntegrationConfig
     """
@@ -271,7 +271,7 @@ def _store_response_headers(headers, span, integration_config):
     :param headers: All the response's http headers, will be filtered through the whitelist
     :type headers: dict or list
     :param span: The Span instance where tags will be stored
-    :type span: ddtrace.Span
+    :type span: ddtrace._trace.span.Span
     :param integration_config: An integration specific config object.
     :type integration_config: ddtrace.settings.IntegrationConfig
     """

ddtrace/contrib/internal/unittest/patch.py

@@ -129,7 +129,9 @@ def _is_marked_as_unskippable(test_object) -> bool:
     )
 
 
-def _update_skipped_elements_and_set_tags(test_module_span: ddtrace.Span, test_session_span: ddtrace.Span):
+def _update_skipped_elements_and_set_tags(
+    test_module_span: ddtrace._trace.span.Span, test_session_span: ddtrace._trace.span.Span
+):
     global _global_skipped_elements
     _global_skipped_elements += 1
 
@@ -141,7 +143,7 @@ def _update_skipped_elements_and_set_tags(test_module_span: ddtrace.Span, test_s
     test_session_span.set_tag_str(test.ITR_DD_CI_ITR_TESTS_SKIPPED, "true")
 
 
-def _store_test_span(item, span: ddtrace.Span):
+def _store_test_span(item, span: ddtrace._trace.span.Span):
     """Store datadog span at `unittest` test instance."""
     item._datadog_span = span
 
@@ -172,7 +174,7 @@ def _is_test(item) -> bool:
     return True
 
 
-def _extract_span(item) -> Union[ddtrace.Span, None]:
+def _extract_span(item) -> Union[ddtrace._trace.span.Span, None]:
     return getattr(item, "_datadog_span", None)
 
 
@@ -187,23 +189,23 @@ def _extract_test_method_name(test_object) -> str:
     return getattr(test_object, "_testMethodName", "")
 
 
-def _extract_session_span() -> Union[ddtrace.Span, None]:
+def _extract_session_span() -> Union[ddtrace._trace.span.Span, None]:
     return getattr(_CIVisibility, "_datadog_session_span", None)
 
 
-def _extract_module_span(module_identifier: str) -> Union[ddtrace.Span, None]:
+def _extract_module_span(module_identifier: str) -> Union[ddtrace._trace.span.Span, None]:
     if hasattr(_CIVisibility, "_unittest_data") and module_identifier in _CIVisibility._unittest_data["modules"]:
         return _CIVisibility._unittest_data["modules"][module_identifier].get("module_span")
     return None
 
 
-def _extract_suite_span(suite_identifier: str) -> Union[ddtrace.Span, None]:
+def _extract_suite_span(suite_identifier: str) -> Union[ddtrace._trace.span.Span, None]:
     if hasattr(_CIVisibility, "_unittest_data") and suite_identifier in _CIVisibility._unittest_data["suites"]:
         return _CIVisibility._unittest_data["suites"][suite_identifier].get("suite_span")
     return None
 
 
-def _update_status_item(item: ddtrace.Span, status: str):
+def _update_status_item(item: ddtrace._trace.span.Span, status: str):
     """
     Sets the status for each Span implementing the test FAIL logic override.
     """
@@ -266,7 +268,7 @@ def _generate_session_resource(test_command: str) -> str:
     return "{}".format(test_command)
 
 
-def _set_test_skipping_tags_to_span(span: ddtrace.Span):
+def _set_test_skipping_tags_to_span(span: ddtrace._trace.span.Span):
     span.set_tag_str(test.ITR_TEST_SKIPPING_ENABLED, "true")
     span.set_tag_str(test.ITR_TEST_SKIPPING_TYPE, TEST)
     span.set_tag_str(test.ITR_TEST_SKIPPING_TESTS_SKIPPED, "false")
@@ -368,7 +370,10 @@ def _finish_remaining_suites_and_modules(seen_suites: dict, seen_modules: dict):
 
 
 def _update_remaining_suites_and_modules(
-    test_module_suite_path: str, test_module_path: str, test_module_span: ddtrace.Span, test_suite_span: ddtrace.Span
+    test_module_suite_path: str,
+    test_module_path: str,
+    test_module_span: ddtrace._trace.span.Span,
+    test_suite_span: ddtrace._trace.span.Span,
 ):
     """
     Updates the remaining test suite and test counter and finishes spans when these have finished their execution.
@@ -384,7 +389,7 @@ def _update_remaining_suites_and_modules(
         _finish_span(test_module_span)
 
 
-def _update_test_skipping_count_span(span: ddtrace.Span):
+def _update_test_skipping_count_span(span: ddtrace._trace.span.Span):
     if _CIVisibility.test_skipping_enabled():
         span.set_metric(test.ITR_TEST_SKIPPING_COUNT, _global_skipped_elements)
 
@@ -622,7 +627,7 @@ def collect_text_test_runner_session(func, instance: unittest.TestSuite, args: t
     return result
 
 
-def _start_test_session_span(instance) -> ddtrace.Span:
+def _start_test_session_span(instance) -> ddtrace._trace.span.Span:
     """
     Starts a test session span and sets the required tags for a `unittest` session instance.
     """
@@ -663,7 +668,7 @@ def _start_test_session_span(instance) -> ddtrace.Span:
     return test_session_span
 
 
-def _start_test_module_span(instance) -> ddtrace.Span:
+def _start_test_module_span(instance) -> ddtrace._trace.span.Span:
     """
     Starts a test module span and sets the required tags for a `unittest` module instance.
     """
@@ -706,7 +711,7 @@ def _start_test_module_span(instance) -> ddtrace.Span:
     return test_module_span
 
 
-def _start_test_suite_span(instance) -> ddtrace.Span:
+def _start_test_suite_span(instance) -> ddtrace._trace.span.Span:
     """
     Starts a test suite span and sets the required tags for a `unittest` suite instance.
     """
@@ -742,7 +747,7 @@ def _start_test_suite_span(instance) -> ddtrace.Span:
     return test_suite_span
 
 
-def _start_test_span(instance, test_suite_span: ddtrace.Span) -> ddtrace.Span:
+def _start_test_span(instance, test_suite_span: ddtrace._trace.span.Span) -> ddtrace._trace.span.Span:
     """
     Starts a test  span and sets the required tags for a `unittest` test instance.
     """
@@ -787,7 +792,7 @@ def _start_test_span(instance, test_suite_span: ddtrace.Span) -> ddtrace.Span:
     return span
 
 
-def _finish_span(current_span: ddtrace.Span):
+def _finish_span(current_span: ddtrace._trace.span.Span):
     """
     Finishes active span and populates span status upwards
     """

ddtrace/internal/ci_visibility/api/_base.py

@@ -76,7 +76,7 @@ class TestVisibilitySessionSettings:
 
     def __post_init__(self):
         if not isinstance(self.tracer, Tracer):
-            raise TypeError("tracer must be a ddtrace.Tracer")
+            raise TypeError("tracer must be a ddtrace._trace.tracer.Tracer")
         if not isinstance(self.workspace_path, Path):
             raise TypeError("root_dir must be a pathlib.Path")
         if not self.workspace_path.is_absolute():

ddtrace/internal/ci_visibility/coverage.py

@@ -131,7 +131,7 @@ def _switch_coverage_context(
 
 
 def _report_coverage_to_span(
-    coverage_data: Coverage, span: ddtrace.Span, root_dir: str, framework: Optional[TEST_FRAMEWORKS] = None
+    coverage_data: Coverage, span: ddtrace._trace.span.Span, root_dir: str, framework: Optional[TEST_FRAMEWORKS] = None
 ):
     # Experimental feature to use internal coverage collection
     if USE_DD_COVERAGE:

ddtrace/internal/ci_visibility/utils.py

@@ -50,7 +50,7 @@ def get_source_lines_for_test_method(
 
 
 def _add_start_end_source_file_path_data_to_span(
-    span: ddtrace.Span, test_method_object, test_name: str, repo_directory: str
+    span: ddtrace._trace.span.Span, test_method_object, test_name: str, repo_directory: str
 ):
     if not test_method_object:
         log.debug(
@@ -75,7 +75,7 @@ def _add_start_end_source_file_path_data_to_span(
         span.set_tag(test.SOURCE_END, end_line)
 
 
-def _add_pct_covered_to_span(coverage_data: dict, span: ddtrace.Span):
+def _add_pct_covered_to_span(coverage_data: dict, span: ddtrace._trace.span.Span):
     if not coverage_data or PCT_COVERED_KEY not in coverage_data:
         log.warning("Tried to add total covered percentage to session span but no data was found")
         return

ddtrace/internal/opentelemetry/trace.py

@@ -103,7 +103,7 @@ def start_span(
         curr_otel_span = get_current_span(context)
         if curr_otel_span is INVALID_SPAN:
             # There is no active datadog/otel span
-            dd_active = None  # type: Optional[Union[ddtrace.trace.Context, ddtrace.Span]]
+            dd_active = None  # type: Optional[Union[ddtrace.trace.Context, ddtrace._trace.span.Span]]
         elif isinstance(curr_otel_span, Span):
             # Get underlying ddtrace span from the active otel span
             dd_active = curr_otel_span._ddspan

ddtrace/internal/runtime/runtime_metrics.py

@@ -78,7 +78,7 @@ def __init__(self, interval=_get_interval_or_default(), tracer=ddtrace.tracer, d
         self._dogstatsd_client: DogStatsd = get_dogstatsd_client(
             self.dogstatsd_url or ddtrace.internal.agent.get_stats_url()
         )
-        self.tracer: Optional[ddtrace.Tracer] = tracer
+        self.tracer: Optional[ddtrace._trace.tracer.Tracer] = tracer
         self._runtime_metrics: RuntimeMetrics = RuntimeMetrics()
         self._services: Set[str] = set()
 
@@ -114,7 +114,7 @@ def _restart(cls):
 
     @classmethod
     def enable(cls, flush_interval=None, tracer=None, dogstatsd_url=None):
-        # type: (Optional[float], Optional[ddtrace.Tracer], Optional[str]) -> None
+        # type: (Optional[float], Optional[ddtrace._trace.tracer.Tracer], Optional[str]) -> None
         with cls._lock:
             if cls._instance is not None:
                 return

ddtrace/internal/runtime/tag_collectors.py

@@ -17,7 +17,7 @@ class RuntimeTagCollector(ValueCollector):
 
 
 class TracerTagCollector(RuntimeTagCollector):
-    """Tag collector for the ddtrace Tracer"""
+    """Tag collector for the ddtrace._trace.tracer.Tracer"""
 
     required_modules = ["ddtrace"]
 

ddtrace/llmobs/_llmobs.py

@@ -298,7 +298,7 @@ def enable(
         api_key: Optional[str] = None,
         env: Optional[str] = None,
         service: Optional[str] = None,
-        _tracer: Optional[ddtrace.Tracer] = None,
+        _tracer: Optional[ddtrace._trace.tracer.Tracer] = None,
     ) -> None:
         """
         Enable LLM Observability tracing.

ddtrace/profiling/collector/stack.pyi

@@ -4,4 +4,4 @@ import ddtrace
 from ddtrace.profiling import collector
 
 class StackCollector(collector.PeriodicCollector):
-    tracer: typing.Optional[ddtrace.Tracer]
+    tracer: typing.Optional[ddtrace._trace.tracer.Tracer]

ddtrace/profiling/exporter/http.py

@@ -40,7 +40,7 @@ class PprofHTTPExporter(pprof.PprofExporter):
 
     def __init__(
         self,
-        tracer: ddtrace.Tracer = ddtrace.tracer,
+        tracer: ddtrace._trace.tracer.Tracer = ddtrace.tracer,
         enable_code_provenance: bool = True,
         api_key: typing.Optional[str] = None,
         timeout: float = config.api_timeout,

ddtrace/runtime/__init__.py

@@ -26,7 +26,7 @@ class RuntimeMetrics(metaclass=_RuntimeMetricsStatus):
 
     @staticmethod
     def enable(tracer=None, dogstatsd_url=None, flush_interval=None):
-        # type: (Optional[ddtrace.Tracer], Optional[str], Optional[float]) -> None
+        # type: (Optional[ddtrace._trace.tracer.Tracer], Optional[str], Optional[float]) -> None
         """
         Enable the runtime metrics collection service.