Repository: pandas-dev/pandas
Pull request: BUG: string methods with NA
- [x] closes #31632 
- [ ] tests added / passed
- [x] passes `black pandas`
- [x] passes `git diff upstream/master -u -- "*.py" | flake8 --diff`
- [ ] whatsnew entry



Problem statement:
Bug in Series.str.repeat with string dtype and sequence of repeats
#### Code Sample, a copy-pastable example if possible

```pytb
In [1]: import pandas as pd

In [2]: s = pd.Series(['a', None], dtype="string")

In [3]: s.str.repeat([1, 2])
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
~/sandbox/pandas/pandas/core/strings.py in rep(x, r)
    781             try:
--> 782                 return bytes.__mul__(x, r)
    783             except TypeError:

TypeError: descriptor '__mul__' requires a 'bytes' object but received a 'NAType'

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
<ipython-input-3-a01827562f7a> in <module>
----> 1 s.str.repeat([1, 2])

~/sandbox/pandas/pandas/core/strings.py in wrapper(self, *args, **kwargs)
   1950                 )
   1951                 raise TypeError(msg)
-> 1952             return func(self, *args, **kwargs)
   1953
   1954         wrapper.__name__ = func_name

~/sandbox/pandas/pandas/core/strings.py in repeat(self, repeats)
   2780     @forbid_nonstring_types(["bytes"])
   2781     def repeat(self, repeats):
-> 2782         result = str_repeat(self._parent, repeats)
   2783         return self._wrap_result(result)
   2784

~/sandbox/pandas/pandas/core/strings.py in str_repeat(arr, repeats)
    785
    786         repeats = np.asarray(repeats, dtype=object)
--> 787         result = libops.vec_binop(com.values_from_object(arr), repeats, rep)
    788         return result
    789

~/sandbox/pandas/pandas/_libs/ops.pyx in pandas._libs.ops.vec_binop()
    239                 result[i] = y
    240             else:
--> 241                 raise
    242
    243     return maybe_convert_bool(result.base)  # `.base` to access np.ndarray

~/sandbox/pandas/pandas/_libs/ops.pyx in pandas._libs.ops.vec_binop()
    232         y = right[i]
    233         try:
--> 234             result[i] = op(x, y)
    235         except TypeError:
    236             if x is None or is_nan(x):

~/sandbox/pandas/pandas/core/strings.py in rep(x, r)
    782                 return bytes.__mul__(x, r)
    783             except TypeError:
--> 784                 return str.__mul__(x, r)
    785
    786         repeats = np.asarray(repeats, dtype=object)

TypeError: descriptor '__mul__' requires a 'str' object but received a 'NAType'
```
#### Problem description

The `str_repeat` method correctly handles NA values when `repeats` is a scalar, but fails when its a sequence.

#### Expected Output

```python
0       a
1    <NA>
dtype: string
```

Hints from issue discussion:
will look into this
Seems to me is_nan() is not able to identify `<NA>`.
In `vec_binop` of ops.pyx if I use `type(x) is pandas._libs.missing.NAType` as an additional condition then the code works. I will look into `is_nan` to fix this.

Commit message: BUG: Series.str.repeat can handle pd.NA for vectored inputs (#31632)

Patch to review:
diff --git a/doc/source/whatsnew/v1.0.2.rst b/doc/source/whatsnew/v1.0.2.rst
index 70aaaa6d0a60d..ac0a3ae42f6db 100644
--- a/doc/source/whatsnew/v1.0.2.rst
+++ b/doc/source/whatsnew/v1.0.2.rst
@@ -29,6 +29,10 @@ Bug fixes
 
 - Using ``pd.NA`` with :meth:`DataFrame.to_json` now correctly outputs a null value instead of an empty object (:issue:`31615`)
 
+**Strings**
+
+- Using ``pd.NA`` with :meth:`Series.str.repeat` now correctly outputs a null value instead of raising error for vector inputs (:issue:`31632`)
+
 .. ---------------------------------------------------------------------------
 
 .. _whatsnew_102.contributors:
diff --git a/pandas/core/strings.py b/pandas/core/strings.py
index 18c7504f2c2f8..3dc964f6c3949 100644
--- a/pandas/core/strings.py
+++ b/pandas/core/strings.py
@@ -778,6 +778,8 @@ def scalar_rep(x):
     else:
 
         def rep(x, r):
+            if isinstance(x, libmissing.NAType):
+                return x
             try:
                 return bytes.__mul__(x, r)
             except TypeError:
diff --git a/pandas/tests/test_strings.py b/pandas/tests/test_strings.py
index 62d26dacde67b..0bc10ac3e7823 100644
--- a/pandas/tests/test_strings.py
+++ b/pandas/tests/test_strings.py
@@ -1153,6 +1153,12 @@ def test_repeat(self):
         assert isinstance(rs, Series)
         tm.assert_series_equal(rs, xp)
 
+        # GH: 31632
+        values = Series(["a", None])
+        result = values.str.repeat([3, 4])
+        exp = Series(["aaa", None])
+        tm.assert_series_equal(result, exp)
+
     def test_match(self):
         # New match behavior introduced in 0.13
         values = Series(["fooBAD__barBAD", np.nan, "foo"])


Code review comments:
diff_hunk: @@ -778,6 +778,8 @@ def scalar_rep(x):
     else:
 
         def rep(x, r):
+            if isinstance(x, libmissing.NAType):
path: pandas/core/strings.py
text: @user1:
do we need the same change for scalar_rep? (and a test that hits it)?

@author:
If sequence passed to repeat has None it is handled...have added a test for it

@user2:
instead of isinstance, just use `if x is NA`

@author:
Thanks for the suggestion. Have updated.