Coverage for /Users/rik/github/cgse/libs/cgse-common/src/egse/decorators.py: 18%

1"""

2 A collection of useful decorator functions.

3"""

4

5 import cProfile

6 import functools

7 import logging

8 import pstats

9 import time

10 import types

11 import warnings

12 from typing import Callable

13 from typing import List

14 from typing import Optional

15

16 import rich

17

18 from egse.settings import Settings

19 from egse.system import get_caller_info

20 from egse.log import logger

21

22

23 def static_vars(**kwargs):

24 """

25 Define static variables in a function.

26

27 The static variable can be accessed with <function name>.<variable name> inside the function body.

28

29 Example:

30 ```python

31 @static_vars(count=0)

32 def special_count():

33 return special_count.count += 2

34 ```

35

36 """

37

38 def decorator(func):

39 for kw in kwargs:

40 setattr(func, kw, kwargs[kw])

41 return func

42

43 return decorator

44

45

46 def implements_protocol(protocol):

47 """

48 Decorator to verify and document protocol compliance at class definition time.

49

50 Usage:

51 @implements_protocol(AsyncRegistryBackend)

52 class MyBackend:

53 ...

54 """

55

56 def decorator(cls):

57 # Add protocol documentation

58 if cls.__doc__:

59 cls.__doc__ += f"\n\nThis class implements the {protocol.__name__} protocol."

60 else:

61 cls.__doc__ = f"This class implements the {protocol.__name__} protocol."

62

63 # Store the protocol for reference

64 cls.__implements_protocol__ = protocol

65

66 # Add runtime verification method

67 def _verify_protocol_compliance(self):

68 if not isinstance(self, protocol):

69 raise TypeError(f"{self.__class__.__name__} does not correctly implement {protocol.__name__}")

70 return True

71

72 cls.verify_protocol_compliance = _verify_protocol_compliance

73

74 # Return the modified class

75 return cls

76

77 return decorator

78

79

80 def dynamic_interface(func) -> Callable:

81 """Adds a static variable `__dynamic_interface` to a method.

82

83 The intended use of this function is as a decorator for functions in an interface class.

84

85 The static variable is currently used by the Proxy class to check if a method

86 is meant to be overridden dynamically. The idea behind this is to loosen the contract

87 of an abstract base class (ABC) into an interface. For an ABC, the abstract methods

88 must be implemented at construction/initialization. This is not possible for the Proxy

89 subclasses as they load their commands (i.e. methods) from the control server, and the

90 method will be added to the Proxy interface after loading. Nevertheless, we like the

91 interface already defined for auto-completion during development or interactive use.

92

93 When a Proxy subclass that implements an interface with methods decorated by

94 the `@dynamic_interface` does overwrite one or more of the decorated methods statically,

95 these methods will not be dynamically overwritten when loading the interface from the

96 control server. A warning will be logged instead.

97 """

98 setattr(func, "__dynamic_interface", True)

99 return func

100

101

102 def query_command(func):

103 """Adds a static variable `__query_command` to a method."""

104

105 setattr(func, "__query_command", True)

106 return func

107

108

109 def transaction_command(func):

110 """Adds a static variable `__transaction_command` to a method."""

111

112 setattr(func, "__transaction_command", True)

113 return func

114

115

116 def read_command(func):

117 """Adds a static variable `__read_command` to a method."""

118

119 setattr(func, "__read_command", True)

120 return func

121

122

123 def write_command(func):

124 """Adds a static variable `__write_command` to a method."""

125

126 setattr(func, "__write_command", True)

127 return func

128

129

130 def average_time(*, name: str = "average_time", level: int = logging.INFO, precision: int = 6) -> Callable:

131 """

132 This is a decorator that is intended mainly as a development aid. When you decorate your function with

133 `@average_time`, the execution time of your function will be kept and accumulated. At anytime in your code,

134 you can request the total execution time and the number of calls:

135

136 @average_time()

137 def my_function():

138 ...

139 total_execution_time, call_count = my_function.report()

140

141 Requesting the report will automatically log the average runtime and the number of calls.

142 If you need to reset the execution time and the number of calls during your testing, use:

143

144 my_function.reset()

145

146 Args:

147 name: A name for the timer that will be used during reporting, default='average_time'

148 level: the required log level, default=logging.INFO

149 precision: the precision used to report the average time, default=6

150

151 Returns:

152 The decorated function.

153

154 """

155

156 def actual_decorator(func):

157 func._run_time = 0.0

158 func._call_count = 0

159

160 def _report_average_time():

161 if func._call_count:

162 average_time = func._run_time / func._call_count

163 logger.log(

164 level,

165 f"{name}: "

166 f"average runtime of {func.__name__! r} is {average_time:.{precision}f}s, "

167 f"#calls = {func._call_count}.",

168 )

169 else:

170 logger.log(level, f"{name}: function {func.__name__! r} was never called.")

171

172 return func._run_time, func._call_count

173

174 def _reset():

175 func._run_time = 0

176 func._call_count = 0

177

178 func.report = _report_average_time

179 func.reset = _reset

180

181 @functools.wraps(func)

182 def wrapper(*args, **kwargs):

183 start_time = time.perf_counter()

184 result = func(*args, **kwargs)

185 end_time = time.perf_counter()

186 func._run_time += end_time - start_time

187 func._call_count += 1

188 return result

189

190 return wrapper

191

192 return actual_decorator

193

194

195 def timer(*, name: str = "timer", level: int = logging.INFO, precision: int = 4):

196 """

197 Print the runtime of the decorated function.

198

199 Args:

200 name: a name for the Timer, will be printed in the logging message

201 level: the logging level for the time message [default=INFO]

202 precision: the number of decimals for the time [default=3 (ms)]

203 """

204

205 def actual_decorator(func):

206 @functools.wraps(func)

207 def wrapper_timer(*args, **kwargs):

208 start_time = time.perf_counter()

209 value = func(*args, **kwargs)

210 end_time = time.perf_counter()

211 run_time = end_time - start_time

212 logger.log(level, f"{name}: Finished {func.__name__! r} in {run_time:.{precision}f} secs")

213 return value

214

215 return wrapper_timer

216

217 return actual_decorator

218

219

220 def async_timer(*, name: str = "timer", level: int = logging.INFO, precision: int = 4):

221 """

222 Print the runtime of the decorated async function.

223

224 Args:

225 name: a name for the Timer, will be printed in the logging message

226 level: the logging level for the time message [default=INFO]

227 precision: the number of decimals for the time [default=3 (ms)]

228 """

229

230 def actual_decorator(func):

231 @functools.wraps(func)

232 async def wrapper_timer(*args, **kwargs):

233 start_time = time.perf_counter()

234 value = await func(*args, **kwargs)

235 end_time = time.perf_counter()

236 run_time = end_time - start_time

237 logger.log(level, f"{name}: Finished {func.__name__! r} in {run_time:.{precision}f} secs")

238 return value

239

240 return wrapper_timer

241

242 return actual_decorator

243

244

245 def time_it(count: int = 1000, precision: int = 4) -> Callable:

246 """Print the runtime of the decorated function.

247

248 This is a simple replacement for the builtin ``timeit`` function. The purpose is to simplify

249 calling a function with some parameters.

250

251 The intended way to call this is as a function:

252

253 value = function(args)

254

255 value = time_it(10_000)(function)(args)

256

257 The `time_it` function can be called as a decorator in which case it will always call the

258 function `count` times which is probably not what you want.

259

260 Args:

261 count (int): the number of executions [default=1000].

262 precision (int): the number of significant digits [default=4]

263

264 Returns:

265 value: the return value of the last function execution.

266

267 See also:

268 the ``Timer`` context manager located in ``egse.system``.

269

270 Usage:

271 ```python

272 @time_it(count=10000)

273 def function(args):

274 pass

275

276 time_it(10000)(function)(args)

277 ```

278 """

279

280 def actual_decorator(func):

281 @functools.wraps(func)

282 def wrapper_timer(*args, **kwargs):

283 value = None

284 start_time = time.perf_counter()

285 for _ in range(count):

286 value = func(*args, **kwargs)

287 end_time = time.perf_counter()

288 run_time = end_time - start_time

289 logging.info(

290 f"Finished {func.__name__! r} in {run_time / count:.{precision}f} secs "

291 f"(total time: {run_time:.2f}s, count: {count})"

292 )

293 return value

294

295 return wrapper_timer

296

297 return actual_decorator

298

299

300 def debug(func):

301 """Logs the function signature and return value."""

302

303 @functools.wraps(func)

304 def wrapper_debug(*args, **kwargs):

305 if __debug__:

306 args_repr = [repr(a) for a in args]

307 kwargs_repr = [f"{k}={v! r}" for k, v in kwargs.items()]

308 signature = ", ".join(args_repr + kwargs_repr)

309 logger.debug(f"Calling {func.__name__}({signature})")

310 value = func(*args, **kwargs)

311 logger.debug(f"{func.__name__! r} returned {value! r}")

312 else:

313 value = func(*args, **kwargs)

314 return value

315

316 return wrapper_debug

317

318

319 def profile_func(

320 output_file: str = None, sort_by: str = "cumulative", lines_to_print: int = None, strip_dirs: bool = False

321) -> Callable:

322 """A time profiler decorator.

323

324 Args:

325 output_file: str or None. Default is None

326 Path of the output file. If only name of the file is given, it's

327 saved in the current directory.

328 If it's None, the name of the decorated function is used.

329 sort_by: str or SortKey enum or tuple/list of str/SortKey enum

330 Sorting criteria for the Stats object.

331 For a list of valid string and SortKey refer to:

332 https://docs.python.org/3/library/profile.html#pstats.Stats.sort_stats

333 lines_to_print: int or None

334 Number of lines to print. Default (None) is for all the lines.

335 This is useful in reducing the size of the printout, especially

336 that sorting by 'cumulative', the time consuming operations

337 are printed toward the top of the file.

338 strip_dirs: bool

339 Whether to remove the leading path info from file names.

340 This is also useful in reducing the size of the printout

341

342 Returns:

343 Profile of the decorated function

344

345 Note:

346 This code was taken from this gist: [a profile

347 decorator](https://gist.github.com/ekhoda/2de44cf60d29ce24ad29758ce8635b78).

348

349 Inspired by and modified the profile decorator of Giampaolo Rodola:

350 [profile decorato](http://code.activestate.com/recipes/577817-profile-decorator/).

351

352

353 """

354

355 def inner(func):

356 @functools.wraps(func)

357 def wrapper(*args, **kwargs):

358 _output_file = output_file or func.__name__ + ".prof"

359 pr = cProfile.Profile()

360 pr.enable()

361 retval = func(*args, **kwargs)

362 pr.disable()

363 pr.dump_stats(_output_file)

364

365 with open(_output_file, "w") as f:

366 ps = pstats.Stats(pr, stream=f)

367 if strip_dirs:

368 ps.strip_dirs()

369 if isinstance(sort_by, (tuple, list)):

370 ps.sort_stats(*sort_by)

371 else:

372 ps.sort_stats(sort_by)

373 ps.print_stats(lines_to_print)

374 return retval

375

376 return wrapper

377

378 return inner

379

380

381 def profile(func):

382 """

383 Prints the function signature and return value to stdout.

384

385 This function checks the `Settings.profiling()` value and only prints out

386 profiling information if this returns True.

387

388 Profiling can be activated with `Settings.set_profiling(True)`.

389 """

390 if not hasattr(profile, "counter"):

391 profile.counter = 0

392

393 @functools.wraps(func)

394 def wrapper_profile(*args, **kwargs):

395 if Settings.profiling():

396 profile.counter += 1

397 args_repr = [repr(a) for a in args]

398 kwargs_repr = [f"{k}={v! r}" for k, v in kwargs.items()]

399 signature = ", ".join(args_repr + kwargs_repr)

400 caller = get_caller_info(level=2)

401 prefix = f"PROFILE[{profile.counter}]: "

402 rich.print(f"{prefix}Calling {func.__name__}({signature})")

403 rich.print(f"{prefix} from {caller.filename} at {caller.lineno}.")

404 value = func(*args, **kwargs)

405 rich.print(f"{prefix}{func.__name__! r} returned {value! r}")

406 profile.counter -= 1

407 else:

408 value = func(*args, **kwargs)

409 return value

410

411 return wrapper_profile

412

413

414 class Profiler:

415 """

416 A simple profiler class that provides some useful functions to profile a function.

417

418 - count: count the number of times this function is executed

419 - duration: measure the total and average duration of the function [seconds]

420

421 Examples:

422 >>> from egse.decorators import Profiler

423 >>> @Profiler.count()

424 ... def square(x):

425 ... return x**2

426

427 >>> x = [square(x) for x in range(1_000_000)]

428

429 >>> print(f"Function 'square' called {square.get_count()} times.")

430 >>> print(square)

431

432 >>> @Profiler.duration()

433 ... def square(x):

434 ... time.sleep(0.1)

435 ... return x**2

436

437 >>> x = [square(x) for x in range(100)]

438

439 >>> print(f"Function 'square' takes on average {square.get_average_duration():.6f} seconds.")

440 >>> print(square)

441

442 """

443

444 class CountCalls:

445 def __init__(self, func):

446 self.func = func

447 self.count = 0

448

449 def __call__(self, *args, **kwargs):

450 self.count += 1

451 return self.func(*args, **kwargs)

452

453 def get_count(self):

454 return self.count

455

456 def reset(self):

457 self.count = 0

458

459 def __str__(self):

460 return f"Function '{self.func.__name__}' was called {self.count} times."

461

462 # The __get__ method is here to make the decorator work with instance methods (methods inside a class)

463 # as well. It ensures that when the decorated method is called on an instance, the self argument is

464 # correctly passed to the method.

465

466 def __get__(self, instance, owner):

467 if instance is None:

468 return self

469 else:

470 return types.MethodType(self, instance)

471

472 class Duration:

473 def __init__(self, func):

474 self.func = func

475 self.duration = 0

476 self.count = 0

477

478 def __call__(self, *args, **kwargs):

479 start = time.perf_counter_ns()

480 response = self.func(*args, **kwargs)

481 self.count += 1

482 self.duration += time.perf_counter_ns() - start

483 return response

484

485 def get_count(self):

486 return self.count

487

488 def get_duration(self):

489 return self.duration / 1_000_000_000

490

491 def get_average_duration(self):

492 return self.duration / 1_000_000_000 / self.count if self.count else 0.0

493

494 def reset(self):

495 self.duration = 0

496 self.count = 0

497

498 def __str__(self):

499 return f"Function '{self.func.__name__}' takes on average {self.get_average_duration():.6f} seconds."

500

501 # The __get__ method is here to make the decorator work with instance methods (methods inside a class)

502 # as well. It ensures that when the decorated method is called on an instance, the self argument is

503 # correctly passed to the method.

504

505 def __get__(self, instance, owner):

506 if instance is None:

507 return self

508 else:

509 return types.MethodType(self, instance)

510

511 @classmethod

512 def count(cls):

513 return cls.CountCalls

514

515 @classmethod

516 def duration(cls):

517 return cls.Duration

518

519

520 def to_be_implemented(func):

521 """Print a warning message that this function/method has to be implemented."""

522

523 @functools.wraps(func)

524 def wrapper_tbi(*args, **kwargs):

525 logger.warning(f"The function/method {func.__name__} is not yet implemented.")

526 return func(*args, **kwargs)

527

528 return wrapper_tbi

529

530

531# Taken and adapted from https://github.com/QCoDeS/Qcodes

532

533

534 def deprecate(reason: Optional[str] = None, alternative: Optional[str] = None) -> Callable:

535 """

536 Deprecate a function or method. This will print a warning with the function name and where

537 it is called from. If the optional parameters `reason` and `alternative` are given, that

538 information will be printed with the warning.

539

540 Examples:

541

542 @deprecate(reason="it doesn't follow PEP8", alternative="set_color()")

543 def setColor(self, color):

544 self.set_color(color)

545

546 Args:

547 reason: provide a short explanation why this function is deprecated. Generates 'because {reason}'

548 alternative: provides an alternative function/parameters to be used. Generates 'Use {alternative}

549 as an alternative'

550

551 Returns:

552 The decorated function.

553 """

554

555 def actual_decorator(func: Callable) -> Callable:

556 @functools.wraps(func)

557 def decorated_func(*args, **kwargs):

558 caller = get_caller_info(2)

559 msg = f'The function "{func.__name__}" used at {caller.filename}:{caller.lineno} is deprecated'

560 if reason is not None:

561 msg += f", because {reason}"

562 if alternative is not None:

563 msg += f". Use {alternative} as an alternative"

564 msg += "."

565 warnings.warn(msg, DeprecationWarning, stacklevel=2)

566 return func(*args, **kwargs)

567

568 decorated_func.__doc__ = (

569 f"This function is DEPRECATED, because {reason}, use {alternative} as an alternative.\n"

570 )

571 return decorated_func

572

573 return actual_decorator

574

575

576 def singleton(cls):

577 """

578 Use class as a singleton.

579

580 from:

581 [Decorator library: Signleton](https://wiki.python.org/moin/PythonDecoratorLibrary#Singleton)

582 """

583

584 cls.__new_original__ = cls.__new__

585

586 @functools.wraps(cls.__new__)

587 def singleton_new(cls, *args, **kw):

588 it = cls.__dict__.get("__it__")

589 if it is not None:

590 return it

591

592 cls.__it__ = it = cls.__new_original__(cls, *args, **kw)

593 it.__init_original__(*args, **kw)

594 return it

595

596 cls.__new__ = singleton_new

597 cls.__init_original__ = cls.__init__

598 cls.__init__ = object.__init__

599

600 return cls

601

602

603 def borg(cls):

604 """

605 Use the Borg pattern to make a class with a shared state between its instances and subclasses.

606

607 from:

608 [we don't need no singleton](

609 http://code.activestate.com/recipes/66531-singleton-we-dont-need-no-stinkin-singleton-the-bo/)

610 """

611

612 cls._shared_state = {}

613 orig_init = cls.__init__

614

615 def new_init(self, *args, **kwargs):

616 self.__dict__ = cls._shared_state

617 orig_init(self, *args, **kwargs)

618

619 cls.__init__ = new_init

620

621 return cls

622

623

624 class classproperty:

625 """Defines a read-only class property.

626

627 Examples:

628

629 >>> class Message:

630 ... def __init__(self, msg):

631 ... self._msg = msg

632 ...

633 ... @classproperty

634 ... def name(cls):

635 ... return cls.__name__

636

637 >>> msg = Message("a simple doctest")

638 >>> assert "Message" == msg.name

639

640 """

641

642 def __init__(self, func):

643 self.func = func

644

645 def __get__(self, instance, owner):

646 return self.func(owner)

647

648 def __set__(self, instance, value):

649 raise AttributeError(

650 f"Cannot change class property '{self.func.__name__}' for class '{instance.__class__.__name__}'."

651 )

652

653

654 class Nothing:

655 """Just to get a nice repr for Nothing. It is kind of a Null object..."""

656

657 def __repr__(self):

658 return "<Nothing>"

659

660

661 def spy_on_attr_change(obj: object, obj_name: str = None) -> None:

662 """

663 Tweak an object to show attributes changing. The changes are reported as WARNING log messages

664 in the `egse.spy` logger.

665

666 Note this is not a decorator, but a function that changes the class of an object.

667

668 Note that this function is a debugging aid and should not be used in production code!

669

670 Args:

671 obj (object): any object that you want to monitor

672 obj_name (str): the variable name of the object that was given in the code, if None than

673 the class name will be printed.

674

675 Example:

676 ```python

677 class X:

678 pass

679

680 x = X()

681 spy_on_attr_change(x, obj_name="x")

682 x.a = 5

683 ```

684

685 From:

686 [Adding a dunder to an object](https://nedbatchelder.com/blog/202206/adding_a_dunder_to_an_object.html)

687 """

688 logger = logging.getLogger("egse.spy")

689

690 class Wrapper(obj.__class__):

691 def __setattr__(self, name, value):

692 old = getattr(self, name, Nothing())

693 logger.warning(f"Spy: in {obj_name or obj.__class__.__name__} -> {name}: {old! r} -> {value! r}")

694 return super().__setattr__(name, value)

695

696 class_name = obj.__class__.__name__

697 obj.__class__ = Wrapper

698 obj.__class__.__name__ = class_name

699

700

701 def retry_with_exponential_backoff(

702 max_attempts: int = 5, initial_wait: float = 1.0, backoff_factor: int = 2, exceptions: List = None

703) -> Callable:

704 """

705 Decorator for retrying a function with exponential backoff.

706

707 This decorator can be applied to a function to handle specified exceptions by

708 retrying the function execution. It will make up to 'max_attempts' attempts with a

709 waiting period that grows exponentially between each attempt (dependent on the backoff_factor).

710 Any exception from the list provided in the `exceptions` argument will be ignored for the

711 given `max_attempts`.

712

713 If after all attempts still an exception is raised, it will be passed through the

714 calling function, otherwise the functions return value will be returned.

715

716 Args:

717 max_attempts: The maximum number of attempts to make.

718 initial_wait: The initial waiting time in seconds before retrying after the first failure.

719 backoff_factor: The factor by which the wait time increases after each failure.

720 exceptions: list of exceptions to ignore, if None all exceptions will be ignored `max_attempts`.

721

722 Returns:

723 The response from the executed function.

724 """

725

726 exceptions = [Exception] if exceptions is None else exceptions

727

728 def actual_decorator(func):

729 @functools.wraps(func)

730 def decorate_func(*args, **kwargs):

731 attempt = 0

732 wait_time = initial_wait

733 last_exception = None

734

735 while attempt < max_attempts:

736 try:

737 response = func(*args, **kwargs) # Attempt to call the function

738 logger.info(f"{func.__name__} successfully executed.")

739 return response

740 except tuple(exceptions) as exc:

741 last_exception = exc

742 attempt += 1

743 logger.info(

744 f"Retry {attempt}: {func.__name__} will be executing again in {wait_time * backoff_factor}s. "

745 f"Received a {last_exception! r}."

746 )

747 time.sleep(wait_time) # Wait before retrying

748 wait_time *= backoff_factor # Increase wait time for the next attempt

749

750 # If the loop completes, all attempts have failed, reraise the last exception

751

752 raise last_exception

753

754 return decorate_func

755

756 return actual_decorator

757

758

759 def retry(times: int = 3, wait: float = 10.0, exceptions: List = None) -> Callable:

760 """

761 Decorator that retries a function multiple times with a delay between attempts.

762

763 This decorator can be applied to a function to handle specified exceptions by

764 retrying the function execution. It will make up to 'times' attempts with a

765 waiting period of 'wait' seconds between each attempt. Any exception from the

766 list provided in the `exceptions` argument will be ignored for the given `times`.

767

768 If after times attempts still an exception is raised, it will be passed through the

769 calling function, otherwise the functions return value will be returned.

770

771 Args:

772 times (int, optional): The number of retry attempts. Defaults to 3.

773 wait (float, optional): The waiting period between retries in seconds. Defaults to 10.0.

774 exceptions (List[Exception] or None, optional): List of exception types to catch and retry.

775 Defaults to None, which catches all exceptions.

776

777 Returns:

778 Callable: The decorated function.

779

780 Example:

781 Apply the retry decorator to a function with specific retry settings:

782

783 ```python

784 @retry(times=5, wait=15.0, exceptions=[ConnectionError, TimeoutError])

785 def my_function():

786 # Function logic here

787 ```

788

789 Note:

790 The decorator catches specified exceptions and retries the function, logging

791 information about each retry attempt.

792

793 """

794

795 exceptions = [Exception] if exceptions is None else exceptions

796

797 def actual_decorator(func: Callable) -> Callable:

798 @functools.wraps(func)

799 def decorated_func(*args, **kwargs):

800 previous_exception = None

801 for n in range(times):

802 try:

803 return func(*args, **kwargs)

804 except tuple(exceptions) as exc:

805 previous_exception = exc

806 if n < times:

807 logger.info(

808 f"Retry {n + 1}: {func.__name__} will be executing again in {wait}s. "

809 f"Received a {previous_exception! r}."

810 )

811 time.sleep(wait)

812 raise previous_exception

813

814 return decorated_func

815

816 return actual_decorator

817

818

819 def execution_count(func):

820 """Counts the number of times the function has been executed."""

821 func._call_count = 0

822

823 def counts():

824 return func._call_count

825

826 def reset():

827 func._call_count = 0

828

829 func.counts = counts

830 func.reset = reset

831

832 @functools.wraps(func)

833 def wrapper(*args, **kwargs):

834 func._call_count += 1

835 value = func(*args, **kwargs)

836 return value

837

838 return wrapper