Repository: numpy/numpy
File: numpy/_core/tests/test_defchararray.py (Python)

Code before:


class TestComparisonsMixed2(TestComparisons):
    """Ticket #1276"""

    def _create_array_a(self):
        return np.array(
            [['abc', 'abcc', '123'],
             ['789', 'abc', 'xyz']], np.str_).view(np.char.chararray)


class TestInformation:
    def _create_array_a(self):
        return np.array([[' abc ', ''],
                           ['12345', 'MixedCase'],
                           ['123 \t 345 \0 ', 'UPPER']]) \
                            .view(np.char.chararray)

    def _create_array_b(self):
        return np.array([[' \u03a3 ', ''],
                           ['12345', 'MixedCase'],
                           ['123 \t 345 \0 ', 'UPPER']]) \
                            .view(np.char.chararray)

    def test_len(self):
        (A, B) = (self._create_array_a(), self._create_array_b())
        assert_(issubclass(np.char.str_len(A).dtype.type, np.integer))
        assert_array_equal(np.char.str_len(A), [[5, 0], [5, 9], [12, 5]])
        assert_array_equal(np.char.str_len(B), [[3, 0], [5, 9], [12, 5]])

    def test_count(self):
        (A, B) = (self._create_array_a(), self._create_array_b())
        assert_(issubclass(A.count('').dtype.type, np.integer))
        assert_array_equal(A.count('a'), [[1, 0], [0, 1], [0, 0]])
        assert_array_equal(A.count('123'), [[0, 0], [1, 0], [1, 0]])
        # Python doesn't seem to like counting NULL characters
        # assert_array_equal(A.count('\0'), [[0, 0], [0, 0], [1, 0]])
        assert_array_equal(A.count('a', 0, 2), [[1, 0], [0, 0], [0, 0]])
        assert_array_equal(B.count('a'), [[0, 0], [0, 1], [0, 0]])
        assert_array_equal(B.count('123'), [[0, 0], [1, 0], [1, 0]])
        # assert_array_equal(B.count('\0'), [[0, 0], [0, 0], [1, 0]])

    def test_endswith(self):
        A = self._create_array_a()
        assert_(issubclass(A.endswith('').dtype.type, np.bool))
        assert_array_equal(A.endswith(' '), [[1, 0], [0, 0], [1, 0]])
        assert_array_equal(A.endswith('3', 0, 3), [[0, 0], [1, 0], [1, 0]])

        def fail():
            A.endswith('3', 'fdjk')


Review comment (suggestion):
If a test only uses `A`, you could do `(A, _) = self._create_arrays()` to discard `B` without binding it to a variable. By convention in Python a variable named `_` is used for this purpose.

Code after:
        return np.array(
            [['abc', 'abcc', '123'],
             ['789', 'abc', 'xyz']], np.str_).view(np.char.chararray)

class TestInformation:
    def A(self):
        return np.array([[' abc ', ''],
                         ['12345', 'MixedCase'],
                         ['123 \t 345 \0 ', 'UPPER']]) \
                          .view(np.char.chararray)

    def B(self):
        return np.array([[' \u03a3 ', ''],
                         ['12345', 'MixedCase'],
                         ['123 \t 345 \0 ', 'UPPER']]) \
                          .view(np.char.chararray)

    def test_len(self):
        A, B = self.A(), self.B()
        assert_(issubclass(np.char.str_len(A).dtype.type, np.integer))
        assert_array_equal(np.char.str_len(A), [[5, 0], [5, 9], [12, 5]])
        assert_array_equal(np.char.str_len(B), [[3, 0], [5, 9], [12, 5]])

    def test_count(self):
        A, B = self.A(), self.B()
        assert_(issubclass(A.count('').dtype.type, np.integer))
        assert_array_equal(A.count('a'), [[1, 0], [0, 1], [0, 0]])
        assert_array_equal(A.count('123'), [[0, 0], [1, 0], [1, 0]])
        # Python doesn't seem to like counting NULL characters
        assert_array_equal(A.count('a', 0, 2), [[1, 0], [0, 0], [0, 0]])
        assert_array_equal(B.count('a'), [[0, 0], [0, 1], [0, 0]])
        assert_array_equal(B.count('123'), [[0, 0], [1, 0], [1, 0]])

    def test_endswith(self):
        A = self.A()
        assert_(issubclass(A.endswith('').dtype.type, np.bool))
        assert_array_equal(A.endswith(' '), [[1, 0], [0, 0], [1, 0]])
        assert_array_equal(A.endswith('3', 0, 3), [[0, 0], [1, 0], [1, 0]])

        def fail():
            A.endswith('3', 'fdjk')

        assert_raises(TypeError, fail)

    @pytest.mark.parametrize(
        "dtype, encode",
        [("U", str),
         ("S", lambda x: x.encode('ascii')),
         ])
    def test_find(self, dtype, encode):
        A = self.A().astype(dtype)