import builtins
import functools
import multiprocessing
import unittest


class EmptyIterClass:
    def __len__(self):
        return 0
    def __getitem__(self, i):
        raise StopIteration


def run_in_subprocess(func, *args):
    with multiprocessing.Pool(1) as pool:
        res = pool.apply_async(func, args=args)
        return res.get(timeout=2)

def get_zero():
    return 0

def run_test_iter_reduce(builtin_name, item, sentinel=None):
    it = iter(item) if sentinel is None else iter(item, sentinel)

    # Backup builtins
    builtins_dict = builtins.__dict__
    orig = {"iter": iter, "reversed": reversed}

    class CustomStr:
        def __init__(self, name, iterator):
            self.name = name
            self.iterator = iterator
        def __hash__(self):
            return hash(self.name)
        def __eq__(self, other):
            # Here we exhaust our iterator, possibly changing
            # its `it_seq` pointer to NULL
            # The `__reduce__` call should correctly get
            # the pointers after this call
            list(self.iterator)
            return other == self.name

    # del is required here
    # since only setting will result in
    # both keys existing with a hash collision
    del builtins_dict[builtin_name]
    builtins_dict[CustomStr(builtin_name, it)] = orig[builtin_name]

    return it.__reduce__()

class Tests(unittest.TestCase):
    
    def test_reduce_mutating_builtins_iter(self):
        # This is a reproducer of issue #101765
        # where iter `__reduce__` calls could lead to a segfault or SystemError
        # depending on the order of C argument evaluation, which is undefined

        types = [
            (EmptyIterClass(),),
            (bytes(3),),
            (bytearray(3),),
            ((1, 2, 3),),
            (get_zero, 0),
            (tuple[int],)  # GenericAlias
        ]

        run_iter = functools.partial(run_test_iter_reduce, "iter")
        # The returned value of `__reduce__` should not only be valid
        # but also *empty*, as `it` was exhausted during `__eq__`
        # i.e "xyz" returns (iter, ("",))

        with self.subTest(case="str"):
            self.assertEqual(
                run_in_subprocess(run_iter, "xyz"),
                (iter, ("",))
            )
        with self.subTest(case="list"):
            self.assertEqual(
                run_in_subprocess(run_iter, [1, 2, 3]),
                (iter, ([],))
            )

        # _PyEval_GetBuiltin is also called for `reversed` in a branch of
        # listiter_reduce_general
        with self.subTest(case="reversed list"):
            self.assertEqual(
                run_in_subprocess(
                    run_test_iter_reduce, 
                    "reversed",
                    reversed([*range(8)])
                ),
                (iter, ([],))
            )

        for case in types:
            with self.subTest(case=case):
                self.assertEqual(
                    run_in_subprocess(run_iter, *case),
                    (iter, ((),))
                )
            

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