mypy fixes (Azure#6641)

SuyogSoti · web-flow · commit f9ab732fe17a · 2019-08-02T15:44:29.000-06:00
* mypy fixes

* use six
diff --git a/sdk/core/azure-core/azure/core/pipeline/policies/distributed_tracing.py b/sdk/core/azure-core/azure/core/pipeline/policies/distributed_tracing.py
@@ -25,47 +25,45 @@
 # --------------------------------------------------------------------------
 """Traces network calls using the implementation library from the settings."""
 
-from azure.core.pipeline import PipelineRequest, PipelineResponse
 from azure.core.tracing.context import tracing_context
-from azure.core.tracing.abstract_span import AbstractSpan
 from azure.core.tracing.common import set_span_contexts
 from azure.core.pipeline.policies import SansIOHTTPPolicy
 from azure.core.settings import settings
 
-try:
-    from urllib.parse import urlparse
-except ImportError:
-    from urlparse import urlparse
+from six.moves import urllib
 
 try:
     from typing import TYPE_CHECKING
 except ImportError:
     TYPE_CHECKING = False
 
 if TYPE_CHECKING:
-    from typing import Any, Optional
-    from azure.core.pipeline.transport import HttpRequest, HttpResponse
+    # the HttpRequest and HttpResponse related type ignores stem from this issue: #5796
+    from azure.core.pipeline.transport import HttpRequest, HttpResponse  # pylint: disable=ungrouped-imports
+    from azure.core.tracing.abstract_span import AbstractSpan  # pylint: disable=ungrouped-imports
+    from azure.core.pipeline import PipelineRequest, PipelineResponse  # pylint: disable=ungrouped-imports
+    from typing import Any, Optional, Dict, List, Union
 
 
 class DistributedTracingPolicy(SansIOHTTPPolicy):
     """The policy to create spans for Azure Calls"""
 
     def __init__(self):
         # type: () -> None
-        self.parent_span_dict = {}
+        self.parent_span_dict = {}  # type: Dict[AbstractSpan, List[Union[AbstractSpan, Any]]]
         self._request_id = "x-ms-client-request-id"
         self._response_id = "x-ms-request-id"
 
     def set_header(self, request, span):
-        # type: (PipelineRequest[HttpRequest], Any) -> None
+        # type: (PipelineRequest, Any) -> None
         """
         Sets the header information on the span.
         """
         headers = span.to_header()
-        request.http_request.headers.update(headers)
+        request.http_request.headers.update(headers)  # type: ignore
 
     def on_request(self, request):
-        # type: (PipelineRequest[HttpRequest], Any) -> None
+        # type: (PipelineRequest) -> None
         parent_span = tracing_context.current_span.get()
         wrapper_class = settings.tracing_implementation()
         original_context = [parent_span, None]
@@ -77,7 +75,7 @@ def on_request(self, request):
         if parent_span is None:
             return
 
-        path = urlparse(request.http_request.url).path
+        path = urllib.parse.urlparse(request.http_request.url).path  # type: ignore
         if not path:
             path = "/"
         child = parent_span.span(name=path)
@@ -99,12 +97,15 @@ def end_span(self, request, response=None):
             if response and self._response_id in response.headers:
                 span.add_attribute(self._response_id, response.headers[self._response_id])
             span.finish()
-            set_span_contexts(*self.parent_span_dict.pop(span, None))
+            original_context = self.parent_span_dict.pop(span, None)
+            if original_context:
+                set_span_contexts(original_context[0], original_context[1])
 
     def on_response(self, request, response):
-        # type: (PipelineRequest[HttpRequest], PipelineResponse[HttpRequest, HttpResponse], Any) -> None
-        self.end_span(request.http_request, response=response.http_response)
+        # type: (PipelineRequest, PipelineResponse) -> None
+        self.end_span(request.http_request, response=response.http_response)  # type: ignore
 
     def on_exception(self, _request):  # pylint: disable=unused-argument
-        # type: (PipelineRequest[HttpRequest], Any) -> bool
-        self.end_span(_request.http_request)
+        # type: (PipelineRequest) -> bool
+        self.end_span(_request.http_request)  # type: ignore
+        return False
diff --git a/sdk/core/azure-core/azure/core/tracing/context.py b/sdk/core/azure-core/azure/core/tracing/context.py
@@ -123,7 +123,7 @@ def __init__(self):
             self.current_span = _ThreadLocalContext("current_span", None)
 
     def with_current_context(self, func):
-        # type: (Callable[[Any], Any]) -> Any
+        # type: (Callable) -> Any
         """
         Passes the current spans to the new context the function will be run in.
         :param func: The function that will be run in the new context
