brain_numpy_utils.py 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091
  1. # Copyright (c) 2019-2021 hippo91 <guillaume.peillex@gmail.com>
  2. # Copyright (c) 2019-2020 Claudiu Popa <pcmanticore@gmail.com>
  3. # Copyright (c) 2021 Pierre Sassoulas <pierre.sassoulas@gmail.com>
  4. # Copyright (c) 2021 Nick Drozd <nicholasdrozd@gmail.com>
  5. # Copyright (c) 2021 Marc Mueller <30130371+cdce8p@users.noreply.github.com>
  6. # Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
  7. # For details: https://github.com/PyCQA/astroid/blob/main/LICENSE
  8. """Different utilities for the numpy brains"""
  9. from typing import Tuple
  10. from astroid.builder import extract_node
  11. from astroid.nodes.node_classes import Attribute, Import, Name, NodeNG
  12. # Class subscript is available in numpy starting with version 1.20.0
  13. NUMPY_VERSION_TYPE_HINTS_SUPPORT = ("1", "20", "0")
  14. def numpy_supports_type_hints() -> bool:
  15. """
  16. Returns True if numpy supports type hints
  17. """
  18. np_ver = _get_numpy_version()
  19. return np_ver and np_ver > NUMPY_VERSION_TYPE_HINTS_SUPPORT
  20. def _get_numpy_version() -> Tuple[str, str, str]:
  21. """
  22. Return the numpy version number if numpy can be imported. Otherwise returns
  23. ('0', '0', '0')
  24. """
  25. try:
  26. import numpy # pylint: disable=import-outside-toplevel
  27. return tuple(numpy.version.version.split("."))
  28. except ImportError:
  29. return ("0", "0", "0")
  30. def infer_numpy_member(src, node, context=None):
  31. node = extract_node(src)
  32. return node.infer(context=context)
  33. def _is_a_numpy_module(node: Name) -> bool:
  34. """
  35. Returns True if the node is a representation of a numpy module.
  36. For example in :
  37. import numpy as np
  38. x = np.linspace(1, 2)
  39. The node <Name.np> is a representation of the numpy module.
  40. :param node: node to test
  41. :return: True if the node is a representation of the numpy module.
  42. """
  43. module_nickname = node.name
  44. potential_import_target = [
  45. x for x in node.lookup(module_nickname)[1] if isinstance(x, Import)
  46. ]
  47. return any(
  48. ("numpy", module_nickname) in target.names or ("numpy", None) in target.names
  49. for target in potential_import_target
  50. )
  51. def looks_like_numpy_member(member_name: str, node: NodeNG) -> bool:
  52. """
  53. Returns True if the node is a member of numpy whose
  54. name is member_name.
  55. :param member_name: name of the member
  56. :param node: node to test
  57. :return: True if the node is a member of numpy
  58. """
  59. if (
  60. isinstance(node, Attribute)
  61. and node.attrname == member_name
  62. and isinstance(node.expr, Name)
  63. and _is_a_numpy_module(node.expr)
  64. ):
  65. return True
  66. if (
  67. isinstance(node, Name)
  68. and node.name == member_name
  69. and node.root().name.startswith("numpy")
  70. ):
  71. return True
  72. return False