brain_hypothesis.py 1.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. # Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
  2. # For details: https://github.com/PyCQA/astroid/blob/main/LICENSE
  3. # Copyright (c) https://github.com/PyCQA/astroid/blob/main/CONTRIBUTORS.txt
  4. """
  5. Astroid hook for the Hypothesis library.
  6. Without this hook pylint reports no-value-for-parameter for use of strategies
  7. defined using the `@hypothesis.strategies.composite` decorator. For example:
  8. from hypothesis import strategies as st
  9. @st.composite
  10. def a_strategy(draw):
  11. return draw(st.integers())
  12. a_strategy()
  13. """
  14. from astroid.manager import AstroidManager
  15. from astroid.nodes.scoped_nodes import FunctionDef
  16. COMPOSITE_NAMES = (
  17. "composite",
  18. "st.composite",
  19. "strategies.composite",
  20. "hypothesis.strategies.composite",
  21. )
  22. def is_decorated_with_st_composite(node) -> bool:
  23. """Return whether a decorated node has @st.composite applied."""
  24. if node.decorators and node.args.args and node.args.args[0].name == "draw":
  25. for decorator_attribute in node.decorators.nodes:
  26. if decorator_attribute.as_string() in COMPOSITE_NAMES:
  27. return True
  28. return False
  29. def remove_draw_parameter_from_composite_strategy(node):
  30. """Given that the FunctionDef is decorated with @st.composite, remove the
  31. first argument (`draw`) - it's always supplied by Hypothesis so we don't
  32. need to emit the no-value-for-parameter lint.
  33. """
  34. del node.args.args[0]
  35. del node.args.annotations[0]
  36. del node.args.type_comment_args[0]
  37. return node
  38. AstroidManager().register_transform(
  39. node_class=FunctionDef,
  40. transform=remove_draw_parameter_from_composite_strategy,
  41. predicate=is_decorated_with_st_composite,
  42. )