annotations.py
1 import inspect 2 import re 3 import types 4 import warnings 5 from functools import wraps 6 from typing import Callable, ParamSpec, TypeVar, overload 7 8 9 def _get_min_indent_of_docstring(docstring_str: str) -> str: 10 """ 11 Get the minimum indentation string of a docstring, based on the assumption 12 that the closing triple quote for multiline comments must be on a new line. 13 Note that based on ruff rule D209, the closing triple quote for multiline 14 comments must be on a new line. 15 16 Args: 17 docstring_str: string with docstring 18 19 Returns: 20 Whitespace corresponding to the indent of a docstring. 21 """ 22 23 if not docstring_str or "\n" not in docstring_str: 24 return "" 25 26 return re.match(r"^\s*", docstring_str.rsplit("\n", 1)[-1]).group() 27 28 29 P = ParamSpec("P") 30 R = TypeVar("R") 31 32 33 @overload 34 def experimental( 35 f: Callable[P, R], 36 version: str | None = None, 37 *, 38 skip: bool = False, 39 ) -> Callable[P, R]: ... 40 41 42 @overload 43 def experimental( 44 f: None = None, 45 version: str | None = None, 46 *, 47 skip: bool = False, 48 ) -> Callable[[Callable[P, R]], Callable[P, R]]: ... 49 50 51 def experimental( 52 f: Callable[P, R] | None = None, 53 version: str | None = None, 54 *, 55 skip: bool = False, 56 ) -> Callable[[Callable[P, R]], Callable[P, R]]: 57 """Decorator / decorator creator for marking APIs experimental in the docstring. 58 59 Args: 60 f: The function to be decorated. 61 version: The version in which the API was introduced as experimental. 62 The version is used to determine whether the API should be considered 63 as stable or not when releasing a new version of MLflow. 64 skip: If True, the automated decorator removal script will skip this 65 decorator. Use this for APIs that are intentionally kept experimental 66 (e.g., because they are still evolving) and should not be auto-promoted 67 to stable. 68 69 Returns: 70 A decorator that adds a note to the docstring of the decorated API, 71 """ 72 if f: 73 return _experimental(f) 74 else: 75 76 def decorator(f: Callable[P, R]) -> Callable[P, R]: 77 return _experimental(f) 78 79 return decorator 80 81 82 def _experimental(api: Callable[P, R]) -> Callable[P, R]: 83 if inspect.isclass(api): 84 api_type = "class" 85 elif inspect.isfunction(api): 86 api_type = "function" 87 elif isinstance(api, (property, types.MethodType)): 88 api_type = "property" 89 else: 90 api_type = str(type(api)) 91 92 indent = _get_min_indent_of_docstring(api.__doc__) if api.__doc__ else "" 93 notice = ( 94 indent + f".. Note:: Experimental: This {api_type} may change or " 95 "be removed in a future release without warning.\n\n" 96 ) 97 if api_type == "property": 98 api.__doc__ = api.__doc__ + "\n\n" + notice if api.__doc__ else notice 99 else: 100 api.__doc__ = notice + api.__doc__ if api.__doc__ else notice 101 return api 102 103 104 def developer_stable(func): 105 """ 106 The API marked here as `@developer_stable` has certain protections associated with future 107 development work. 108 Classes marked with this decorator implicitly apply this status to all methods contained within 109 them. 110 111 APIs that are annotated with this decorator are guaranteed (except in cases of notes below) to: 112 - maintain backwards compatibility such that earlier versions of any MLflow client, cli, or 113 server will not have issues with any changes being made to them from an interface perspective. 114 - maintain a consistent contract with respect to existing named arguments such that 115 modifications will not alter or remove an existing named argument. 116 - maintain implied or declared types of arguments within its signature. 117 - maintain consistent behavior with respect to return types. 118 119 Note: Should an API marked as `@developer_stable` require a modification for enhanced feature 120 functionality, a deprecation warning will be added to the API well in advance of its 121 modification. 122 123 Note: Should an API marked as `@developer_stable` require patching for any security reason, 124 advanced notice is not guaranteed and the labeling of such API as stable will be ignored 125 for the sake of such a security patch. 126 127 """ 128 return func 129 130 131 _DEPRECATED_MARK_ATTR_NAME = "__deprecated" 132 133 134 def mark_deprecated(func): 135 """ 136 Mark a function as deprecated by setting a private attribute on it. 137 """ 138 setattr(func, _DEPRECATED_MARK_ATTR_NAME, True) 139 140 141 def is_marked_deprecated(func): 142 """ 143 Is the function marked as deprecated. 144 """ 145 return getattr(func, _DEPRECATED_MARK_ATTR_NAME, False) 146 147 148 def deprecated(alternative: str | None = None, since: str | None = None, impact: str | None = None): 149 """Annotation decorator for marking APIs as deprecated in docstrings and raising a warning if 150 called. 151 152 Args: 153 alternative: The name of a superseded replacement function, method, 154 or class to use in place of the deprecated one. 155 since: A version designator defining during which release the function, 156 method, or class was marked as deprecated. 157 impact: Indication of whether the method, function, or class will be 158 removed in a future release. 159 160 Returns: 161 Decorated function or class. 162 """ 163 164 def deprecated_decorator(obj): 165 since_str = f" since {since}" if since else "" 166 impact_str = impact or "This method will be removed in a future release." 167 168 qual_name = f"{obj.__module__}.{obj.__qualname__}" 169 notice = f"``{qual_name}`` is deprecated{since_str}. {impact_str}" 170 if alternative and alternative.strip(): 171 notice += f" Use ``{alternative}`` instead." 172 173 if inspect.isclass(obj): 174 original_init = obj.__init__ 175 176 @wraps(original_init) 177 def new_init(self, *args, **kwargs): 178 warnings.warn(notice, category=FutureWarning, stacklevel=2) 179 original_init(self, *args, **kwargs) 180 181 obj.__init__ = new_init 182 183 if obj.__doc__: 184 obj.__doc__ = f".. Warning:: {notice}\n{obj.__doc__}" 185 else: 186 obj.__doc__ = f".. Warning:: {notice}" 187 188 mark_deprecated(obj) 189 return obj 190 191 elif isinstance(obj, (types.FunctionType, types.MethodType)): 192 193 @wraps(obj) 194 def deprecated_func(*args, **kwargs): 195 warnings.warn(notice, category=FutureWarning, stacklevel=2) 196 return obj(*args, **kwargs) 197 198 if obj.__doc__: 199 indent = _get_min_indent_of_docstring(obj.__doc__) 200 deprecated_func.__doc__ = f"{indent}.. Warning:: {notice}\n{obj.__doc__}" 201 else: 202 deprecated_func.__doc__ = f".. Warning:: {notice}" 203 204 mark_deprecated(deprecated_func) 205 return deprecated_func 206 207 else: 208 return obj 209 210 return deprecated_decorator 211 212 213 def deprecated_parameter(old_param: str, new_param: str, version: str | None = None): 214 """ 215 Decorator to handle deprecated parameter renaming with automatic warning and forwarding. 216 217 This decorator: 218 1. Accepts the deprecated parameter in the function signature 219 2. Emits a deprecation warning when the old parameter is used 220 3. Maps the old parameter value to the new parameter name 221 4. Hides the deprecated parameter from documentation 222 5. Keeps the function body untouched 223 224 Args: 225 old_param: The deprecated parameter name 226 new_param: The new parameter name to use instead 227 version: Optional version when the deprecation will be removed 228 229 Example: 230 @deprecated_parameter("request_id", "trace_id", version="4.0.0") 231 def search_traces(trace_id: str | None = None): 232 # Function body uses trace_id directly 233 return trace_id 234 235 # Users can still call with old parameter: 236 search_traces(request_id="123") # Issues warning, forwards to trace_id 237 search_traces(trace_id="123") # No warning 238 """ 239 240 def decorator(func): 241 sig = inspect.signature(func) 242 params = dict(sig.parameters) 243 244 if new_param not in params: 245 raise ValueError( 246 f"New parameter '{new_param}' not found in function '{func.__name__}' signature" 247 ) 248 249 @wraps(func) 250 def wrapper(*args, **kwargs): 251 if old_param in kwargs: 252 old_value = kwargs.pop(old_param) 253 254 version_msg = f" and will be removed in version {version}" if version else "" 255 warnings.warn( 256 f"Parameter '{old_param}' is deprecated{version_msg}. " 257 f"Please use '{new_param}' instead.", 258 category=FutureWarning, 259 stacklevel=2, 260 ) 261 262 if new_param in kwargs: 263 raise ValueError( 264 f"Cannot specify both '{old_param}' (deprecated) and '{new_param}'. " 265 f"Use '{new_param}' only." 266 ) 267 268 kwargs[new_param] = old_value 269 270 return func(*args, **kwargs) 271 272 # Update the wrapper's signature to include the deprecated parameter as keyword-only 273 # but keep it out of documentation 274 wrapper.__signature__ = sig 275 276 # Update docstring to note the deprecation (if docstring exists) 277 if func.__doc__: 278 indent = _get_min_indent_of_docstring(func.__doc__) 279 deprecation_note = ( 280 f"{indent}.. Note:: Parameter ``{old_param}`` is deprecated. " 281 f"Use ``{new_param}`` instead." 282 ) 283 wrapper.__doc__ = f"{deprecation_note}\n{func.__doc__}" 284 285 return wrapper 286 287 return decorator 288 289 290 def keyword_only(func): 291 """A decorator that forces keyword arguments in the wrapped method.""" 292 293 @wraps(func) 294 def wrapper(*args, **kwargs): 295 if len(args) > 0: 296 raise TypeError(f"Method {func.__name__} only takes keyword arguments.") 297 return func(**kwargs) 298 299 indent = _get_min_indent_of_docstring(wrapper.__doc__) if wrapper.__doc__ else "" 300 notice = indent + ".. note:: This method requires all argument be specified by keyword.\n" 301 wrapper.__doc__ = notice + wrapper.__doc__ if wrapper.__doc__ else notice 302 303 return wrapper 304 305 306 def filter_user_warnings_once(func): 307 """A decorator that filter user warnings to only show once in the wrapped method.""" 308 309 @wraps(func) 310 def wrapper(*args, **kwargs): 311 with warnings.catch_warnings(): 312 warnings.simplefilter("once", category=UserWarning) 313 return func(*args, **kwargs) 314 315 return wrapper 316 317 318 def requires_sql_backend(func): 319 """ 320 Decorator for marking APIs that require a SQL-based tracking backend. 321 322 This decorator: 323 1. Adds a note to the docstring indicating SQL backend requirement 324 2. When used with FileStore, raises a helpful error message 325 326 The decorator should be applied to methods in AbstractStore that are only 327 implemented in SQL-based backends (SQLAlchemyStore, RestStore). 328 """ 329 indent = _get_min_indent_of_docstring(func.__doc__) if func.__doc__ else "" 330 notice = ( 331 indent + ".. Note:: This method requires a SQL-based tracking backend " 332 "(e.g., SQLite, PostgreSQL, MySQL). It is not supported with FileStore.\n\n" 333 ) 334 func.__doc__ = notice + func.__doc__ if func.__doc__ else notice 335 336 func._requires_sql_backend = True 337 338 return func