Skip to content

Commit

Permalink
Make imag and real attributes numpy.ndarray
Browse files Browse the repository at this point in the history
  • Loading branch information
hippo91 authored Feb 10, 2020
1 parent 3a1b900 commit f8b7797
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 2 deletions.
4 changes: 4 additions & 0 deletions ChangeLog
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,10 @@ What's New in astroid 2.4.0?
Release Date: TBA


* Numpy ``ndarray`` attributes ``imag`` and ``real`` are now inferred as ``ndarray``.

Close PyCQA/pylint#3322

* Added a call to ``register_transform`` for all functions of the ``brain_numpy_core_multiarray``
module in case the current node is an instance of ``astroid.Name``

Expand Down
4 changes: 2 additions & 2 deletions astroid/brain/brain_numpy_ndarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,11 +24,11 @@ def __init__(self, shape, dtype=float, buffer=None, offset=0,
self.dtype = None
self.flags = None
self.flat = None
self.imag = None
self.imag = np.ndarray([0, 0])
self.itemsize = None
self.nbytes = None
self.ndim = None
self.real = None
self.real = np.ndarray([0, 0])
self.shape = numpy.ndarray([0, 0])
self.size = None
self.strides = None
Expand Down
31 changes: 31 additions & 0 deletions tests/unittest_brain_numpy_ndarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,18 @@ def _inferred_ndarray_method_call(self, func_name):
)
return node.infer()

def _inferred_ndarray_attribute(self, attr_name):
node = builder.extract_node(
"""
import numpy as np
test_array = np.ndarray((2, 2))
test_array.{:s}
""".format(
attr_name
)
)
return node.infer()

def test_numpy_function_calls_inferred_as_ndarray(self):
"""
Test that some calls to numpy functions are inferred as numpy.ndarray
Expand All @@ -136,6 +148,25 @@ def test_numpy_function_calls_inferred_as_ndarray(self):
),
)

def test_numpy_ndarray_attribute_inferred_as_ndarray(self):
"""
Test that some numpy ndarray attributes are inferred as numpy.ndarray
"""
licit_array_types = ".ndarray"
for attr_ in ("real", "imag"):
with self.subTest(typ=attr_):
inferred_values = list(self._inferred_ndarray_attribute(attr_))
self.assertTrue(
len(inferred_values) == 1,
msg="Too much inferred value for {:s}".format(attr_),
)
self.assertTrue(
inferred_values[-1].pytype() in licit_array_types,
msg="Illicit type for {:s} ({})".format(
attr_, inferred_values[-1].pytype()
),
)


if __name__ == "__main__":
unittest.main()

0 comments on commit f8b7797

Please sign in to comment.