Support integer in-place multiplication and preserve empty cycle keys

cycler/__init__.py

@@ -77,10 +77,10 @@ def _process_keys(
     keys : set
         The keys in the composition of the two cyclers.
     """
-    l_peek: dict[K, V] = next(iter(left)) if left != [] else {}
-    r_peek: dict[K, V] = next(iter(right)) if right is not None else {}
-    l_key: set[K] = set(l_peek.keys())
-    r_key: set[K] = set(r_peek.keys())
+    l_key: set[K] = left.keys if isinstance(left, Cycler) else set(next(iter(left), {}))
+    r_key: set[K] = set()
+    if right is not None:
+        r_key = right.keys if isinstance(right, Cycler) else set(next(iter(right), {}))
     if common_keys := l_key & r_key:
         raise ValueError(
             f"Cannot compose overlapping cycles, duplicate key(s): {common_keys}"
@@ -166,6 +166,7 @@ def __init__(
             self._left: Cycler[K, V] | list[dict[K, V]] = Cycler(
                 left._left, left._right, left._op
             )
+            self._left._keys = left.keys
         elif left is not None:
             # Need to copy the dictionary or else that will be a residual
             # mutable that could lead to strange errors
@@ -177,6 +178,7 @@ def __init__(
             self._right: Cycler[K, V] | None = Cycler(
                 right._left, right._right, right._op
             )
+            self._right._keys = right.keys
         else:
             self._right = None
 
@@ -369,12 +371,18 @@ def __iadd__(self, other: Cycler[K, V]) -> Cycler[K, V]:  # type: ignore[misc]
 
     def __imul__(self, other: Cycler[K, V] | int) -> Cycler[K, V]:  # type: ignore[misc]
         """
-        In-place outer product of two cyclers (`itertools.product`).
+        In-place outer product of two cyclers or integer multiplication.
 
         Parameters
         ----------
-        other : Cycler
+        other : Cycler or int
         """
+        if isinstance(other, int):
+            multiplied = self * other
+            self._left = multiplied._left
+            self._right = multiplied._right
+            self._op = multiplied._op
+            return self
         if not isinstance(other, Cycler):
             raise TypeError("Cannot *= with a non-Cycler object")
         # True shallow copy of self is fine since this is in-place

test_cycler.py

@@ -80,6 +80,38 @@ def test_inplace_add_len_mismatch():
         c1 += c3
 
 
+@pytest.mark.parametrize('factor', [-2, 0, 1, 2, 3])
+@pytest.mark.parametrize('make_cycle', [
+    lambda: cycler(c='rb'),
+    lambda: cycler(c='rb') + cycler(lw=[1, 2]),
+    lambda: cycler(c='rb') * cycler(lw=[1, 2]),
+], ids=['base', 'sum', 'product'])
+def test_inplace_integer_mul(make_cycle, factor):
+    c = make_cycle()
+    alias = c
+    expected = c * factor
+    c *= factor
+    assert c is alias
+    _cycles_equal(c, expected)
+    assert c.keys == expected.keys
+    assert c.by_key() == expected.by_key()
+    assert repr(c) == repr(expected)
+
+
+@pytest.mark.parametrize('factor', [-2, 0])
+@pytest.mark.parametrize('make_cycle', [
+    lambda: cycler(c='rb'),
+    lambda: cycler(c='rb') + cycler(lw=[1, 2]),
+    lambda: cycler(c='rb') * cycler(lw=[1, 2]),
+], ids=['base', 'sum', 'product'])
+def test_integer_mul_empty_preserves_keys(make_cycle, factor):
+    c = make_cycle()
+    result = c * factor
+    assert list(result) == []
+    assert len(result) == 0
+    assert result.keys == c.keys
+
+
 def test_inplace_mul():
     c3 = cycler(c='rgb')
     c4 = cycler(lw=range(3))