From 990d32b9dfec5f0ebdce8e9e79a61d2a6772c8d6 Mon Sep 17 00:00:00 2001 From: Alexander Dokuchaev Date: Thu, 19 Dec 2024 16:52:13 +0200 Subject: [PATCH 1/3] mypy --- nncf/telemetry/decorator.py | 12 +++++++----- nncf/telemetry/events.py | 8 ++++---- nncf/telemetry/extractors.py | 8 ++++---- nncf/telemetry/wrapper.py | 28 ++++++++++++++-------------- 4 files changed, 29 insertions(+), 27 deletions(-) diff --git a/nncf/telemetry/decorator.py b/nncf/telemetry/decorator.py index 23ee56acc86..f88b04d2c36 100644 --- a/nncf/telemetry/decorator.py +++ b/nncf/telemetry/decorator.py @@ -10,7 +10,7 @@ # limitations under the License. import functools import inspect -from typing import Callable, List, Union +from typing import Any, Callable, List, Optional, TypeVar, Union from nncf.telemetry.events import MODEL_BASED_CATEGORY from nncf.telemetry.events import get_current_category @@ -21,6 +21,8 @@ from nncf.telemetry.extractors import VerbatimTelemetryExtractor from nncf.telemetry.wrapper import telemetry +TFunction = TypeVar("TFunction", bound=Callable[..., Any]) + class tracked_function: """ @@ -29,7 +31,7 @@ class tracked_function: function execution. The category of the session and events will be determined by parameters to the decorator. """ - def __init__(self, category: str = None, extractors: List[Union[str, TelemetryExtractor]] = None): + def __init__(self, category: str = None, extractors: Optional[List[Union[str, TelemetryExtractor]]] = None) -> None: """ :param category: A category to be attributed to the events. If set to None, no events will be sent. :param extractors: Add argument names in this list as string values to send an event with an "action" equal to @@ -44,11 +46,11 @@ def __init__(self, category: str = None, extractors: List[Union[str, TelemetryEx else: self._collectors = [] - def __call__(self, fn: Callable) -> Callable: + def __call__(self, fn: TFunction) -> TFunction: fn_signature = inspect.signature(fn) @functools.wraps(fn) - def wrapped(*args, **kwargs): + def wrapped(*args: Any, **kwargs: Any) -> Any: bound_args = fn_signature.bind(*args, **kwargs) bound_args.apply_defaults() @@ -82,4 +84,4 @@ def wrapped(*args, **kwargs): telemetry.end_session(self._category) return retval - return wrapped + return wrapped # type: ignore[return-value] diff --git a/nncf/telemetry/events.py b/nncf/telemetry/events.py index a048fbbb0c7..b0966a5c77c 100644 --- a/nncf/telemetry/events.py +++ b/nncf/telemetry/events.py @@ -10,7 +10,7 @@ # limitations under the License. from contextlib import contextmanager -from typing import Optional, TypeVar +from typing import Generator, Optional, TypeVar from nncf.common.utils.backend import BackendType from nncf.common.utils.backend import get_backend @@ -25,12 +25,12 @@ # Dynamic categories MODEL_BASED_CATEGORY = "model_based" -CURRENT_CATEGORY = None +CURRENT_CATEGORY: Optional[str] = None TModel = TypeVar("TModel") -def _set_current_category(category: str): +def _set_current_category(category: Optional[str]) -> None: global CURRENT_CATEGORY CURRENT_CATEGORY = category @@ -56,7 +56,7 @@ def get_model_based_category(model: TModel) -> str: @contextmanager -def telemetry_category(category: str) -> str: +def telemetry_category(category: Optional[str]) -> Generator[Optional[str], None, None]: previous_category = get_current_category() _set_current_category(category) yield category diff --git a/nncf/telemetry/extractors.py b/nncf/telemetry/extractors.py index a06d8f8daa5..289830bbf6a 100644 --- a/nncf/telemetry/extractors.py +++ b/nncf/telemetry/extractors.py @@ -14,7 +14,7 @@ from enum import Enum from typing import Any, Optional, Union -SerializableData = Union[str, Enum] +SerializableData = Union[str, Enum, bool] @dataclass @@ -26,8 +26,8 @@ class CollectedEvent: """ name: str - data: SerializableData = None # GA limitations - int_data: int = None + data: Optional[SerializableData] = None # GA limitations + int_data: Optional[int] = None class TelemetryExtractor(ABC): @@ -35,7 +35,7 @@ class TelemetryExtractor(ABC): Interface for custom telemetry extractors, to be used with the `nncf.telemetry.tracked_function` decorator. """ - def __init__(self, argname: Optional[str] = None): + def __init__(self, argname: str = "") -> None: self._argname = argname @property diff --git a/nncf/telemetry/wrapper.py b/nncf/telemetry/wrapper.py index 2cc32af0158..eb150c3e1ee 100644 --- a/nncf/telemetry/wrapper.py +++ b/nncf/telemetry/wrapper.py @@ -13,7 +13,7 @@ import sys from abc import ABC from abc import abstractmethod -from typing import Callable, Optional +from typing import Any, Callable, Optional from unittest.mock import MagicMock from nncf.common.logging import nncf_logger @@ -29,7 +29,7 @@ class ITelemetry(ABC): # https://support.google.com/analytics/answer/1033068 @abstractmethod - def start_session(self, category: str, **kwargs): + def start_session(self, category: str, **kwargs: Any) -> None: """ Sends a message about starting of a new session. @@ -45,9 +45,9 @@ def send_event( event_action: str, event_label: str, event_value: Optional[int] = None, - force_send=False, - **kwargs, - ): + force_send: bool = False, + **kwargs: Any, + ) -> None: """ Send single event. @@ -61,7 +61,7 @@ def send_event( """ @abstractmethod - def end_session(self, category: str, **kwargs): + def end_session(self, category: str, **kwargs: Any) -> None: """ Sends a message about ending of the current session. @@ -78,7 +78,7 @@ def skip_if_raised(func: Callable[..., None]) -> Callable[..., None]: """ @functools.wraps(func) - def wrapped(*args, **kwargs): + def wrapped(*args: Any, **kwargs: Any) -> None: try: func(*args, **kwargs) @@ -91,7 +91,7 @@ def wrapped(*args, **kwargs): class NNCFTelemetry(ITelemetry): MEASUREMENT_ID = "G-W5E9RNLD4H" - def __init__(self): + def __init__(self) -> None: self._app_name = "nncf" self._app_version = __version__ try: @@ -108,7 +108,7 @@ def __init__(self): nncf_logger.debug(f"Failed to instantiate telemetry object: exception {e}") @skip_if_raised - def start_session(self, category: str, **kwargs): + def start_session(self, category: str, **kwargs: Any) -> None: self._impl.start_session(category, **kwargs) @skip_if_raised @@ -118,9 +118,9 @@ def send_event( event_action: str, event_label: str, event_value: Optional[int] = None, - force_send=False, - **kwargs, - ): + force_send: bool = False, + **kwargs: Any, + ) -> None: if event_value is None: event_value = 1 self._impl.send_event( @@ -135,12 +135,12 @@ def send_event( ) @skip_if_raised - def end_session(self, category: str, **kwargs): + def end_session(self, category: str, **kwargs: Any) -> None: self._impl.end_session(category, **kwargs) try: - from openvino_telemetry import Telemetry + from openvino_telemetry import Telemetry # type: ignore telemetry = NNCFTelemetry() except ImportError: From 83c749bfb4e8c22e6c3846e5d33a2a0394125d04 Mon Sep 17 00:00:00 2001 From: Alexander Dokuchaev Date: Thu, 19 Dec 2024 17:32:49 +0200 Subject: [PATCH 2/3] f --- nncf/telemetry/decorator.py | 2 +- nncf/telemetry/extractors.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/nncf/telemetry/decorator.py b/nncf/telemetry/decorator.py index f88b04d2c36..eba2be711fb 100644 --- a/nncf/telemetry/decorator.py +++ b/nncf/telemetry/decorator.py @@ -61,7 +61,7 @@ def wrapped(*args: Any, **kwargs: Any) -> Any: events: List[CollectedEvent] = [] for collector in self._collectors: argname = collector.argname - argvalue = bound_args.arguments[argname] if argname is not None else None + argvalue = bound_args.arguments[argname] if argname else None event = collector.extract(argvalue) events.append(event) diff --git a/nncf/telemetry/extractors.py b/nncf/telemetry/extractors.py index 289830bbf6a..624b0b1f26a 100644 --- a/nncf/telemetry/extractors.py +++ b/nncf/telemetry/extractors.py @@ -35,7 +35,7 @@ class TelemetryExtractor(ABC): Interface for custom telemetry extractors, to be used with the `nncf.telemetry.tracked_function` decorator. """ - def __init__(self, argname: str = "") -> None: + def __init__(self, argname: str = ""): self._argname = argname @property From 7d0f152286a693b4bc858a32fc293c11b8dc39d9 Mon Sep 17 00:00:00 2001 From: Alexander Dokuchaev Date: Thu, 19 Dec 2024 17:40:24 +0200 Subject: [PATCH 3/3] pyproject.toml --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index eb6067af2c9..45344934e07 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,6 +100,7 @@ files = [ "nncf/common/utils/", "nncf/common/tensor_statistics", "nncf/experimental/torch2", + "nncf/telemetry/", ] [tool.ruff]