Coverage for /home/crpier/Projects/snektest/snektest/utils.py: 100%

1from collections.abc import AsyncGenerator, Callable, Generator 

2from inspect import isasyncgen 

3from types import CodeType 

4from typing import Any 

5 

6from snektest.models import Param 

7 

8TEST_ATTR_NAME = "__is_snektest__test__" 

9TEST_ATTR_VALUE = object() 

10 

11PARAMS_ATTR_NAME = "__snektest_params__" 

12MARKERS_ATTR_NAME = "__snektest_markers__" 

13 

14 

15def mark_test_function( 

16 func: Callable[..., Any], 

17 params: tuple[list[Param[Any]], ...], 

18 markers: tuple[str, ...], 

19) -> None: 

20 """Mark a function as a test and store its parameters.""" 

21 setattr(func, TEST_ATTR_NAME, TEST_ATTR_VALUE) 

22 setattr(func, PARAMS_ATTR_NAME, Param.to_dict(params)) 

23 setattr(func, MARKERS_ATTR_NAME, markers) 

24 

25 

26def is_test_function(func: Callable[..., Any]) -> bool: 

27 """Check if a function is marked as a test.""" 

28 return getattr(func, TEST_ATTR_NAME, None) is TEST_ATTR_VALUE 

29 

30 

31def get_test_function_params( 

32 func: Callable[..., Any], 

33) -> dict[str, tuple[Param[Any], ...]]: 

34 """Get the parameters dict for a test function.""" 

35 return getattr(func, PARAMS_ATTR_NAME) 

36 

37 

38def get_test_function_markers(func: Callable[..., Any]) -> tuple[str, ...]: 

39 """Get the markers tuple for a test function.""" 

40 return getattr(func, MARKERS_ATTR_NAME, ()) 

41 

42 

43def get_code_from_generator( 

44 generator: AsyncGenerator[Any] | Generator[Any], 

45) -> CodeType: 

46 """Get the code object from a generator.""" 

47 return generator.ag_code if isasyncgen(generator) else generator.gi_code # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportAttributeAccessIssue, reportUnknownVariableType] 

48 

49 

50def get_func_name_from_generator( 

51 generator: AsyncGenerator[Any] | Generator[Any], 

52) -> str: 

53 """Get the code object from a generator.""" 

54 return ( 

55 generator.ag_code.co_name 

56 if isasyncgen(generator) 

57 else generator.gi_code.co_name # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportAttributeAccessIssue, reportUnknownVariableType] 

58 )