Repository: radiasoft/pykern

Issue:
pkcompat.zip_strict

Discussion:
robnagler: I spent a little time trying to figure out how this magic little function works:

```py
def _zip_strict(iterables):
    """Code from https://stackoverflow.com/a/69485272"""

    def _first_tail():
        nonlocal first_stopped
        first_stopped = True
        return
        yield

    def _zip_tail():
        if not first_stopped:
            raise ValueError("first iterable is longer")
        for _ in itertools.chain.from_iterable(rest):
            raise ValueError("first iterable is shorter")
            yield

    if len(iterables) < 2:
        return zip(*iterables)
    first_stopped = False
    iterables = iter(iterables)
    first = itertools.chain(next(iterables), _first_tail())
    rest = list(map(iter, iterables))
    return itertools.chain(zip(first, *rest), _zip_tail())
```

Some definitions:
- `next(any_iterator)` must return `StopIteration` after it was already raised
- `iter` on an iterator returns self
- `zip` doesn't require an iterator 
- `zip` stops iterating the rest of its arguments once an argument is exhausted. This is similar to `in`, btw.
- `chain` concatenates a sequence of iterables
- `from_iterable` concatenates  a single argument of iterables 

Algorithm:
- `first` has a "watcher" (`_first_tail`) that is an empty generator (forced by the `yield`) 
- `rest` has no watchers
- `zip` has a watcher (`_zip_tail`) that checks:
  -  if `_first_tail` did not execute, then `rest` was shorter than `iterable[0]`
  - Otherwise, check if there's nothing left in `rest`, that is, `next` on each element of `rest` must return `StopIteration`

Notes:
- It's important that `rest` be iterators, not iterables. That was confusing to me. Only iterators report `StopIteration`, not sequences. 
- `zip` does not care if something is a sequence or an iterator. It will call `iter` on all its arguments. 
- I didn't understand why the `map` had to call `iter`: it's to ensure that `from_iterable`, which also calls `iter`, doesn't start from the beginning. 
- The algorithm above is highly efficient, because `iter` returns self on an iterator. It's a no-op, and itertools is written in C.
- All the work is done by zip until the last iteration when `_first_tail` and `_zip_tail` execute once. No matter how big the iterables are, the code will run almost linearly (as shown [in the graph](Code from https://stackoverflow.com/a/69485272)).

HTH.



Patch:
diff --git a/pykern/pkcompat.py b/pykern/pkcompat.py
index d446bc8..69e2baf 100644
--- a/pykern/pkcompat.py
+++ b/pykern/pkcompat.py
@@ -1,18 +1,24 @@
 # -*- coding: utf-8 -*-
-"""Python 2 and 3 compatibility str routines
+"""Backwards and forward compatible Python utilities
 
-:mod:`six` and :mod:`future.utils` do most things, but there are some missing
-things here
+Functions here will be available forever, even when functions are
+removed from Python.
 
-:copyright: Copyright (c) 2015-2022 RadiaSoft LLC.  All Rights Reserved.
+Some functions are no longer necessary, e.g. `unicode_getcwd`. Others
+are still convenient, e.g. `from_bytes` is useful when you don't know
+whether it is a `str` or not.
+
+:copyright: Copyright (c) 2015-2023 RadiaSoft LLC.  All Rights Reserved.
 :license: http://www.apache.org/licenses/LICENSE-2.0.html
+
 """
 # Limit pykern imports so avoid dependency issues for pkconfig
-from __future__ import absolute_import, division, print_function
 import inspect
+import itertools
 import locale
 import os
 import subprocess
+import sys
 
 
 def from_bytes(value):
@@ -90,6 +96,20 @@ def unicode_unescape(value):
     return value.encode("utf-8").decode("unicode-escape")
 
 
+def zip_strict(*iterables):
+    """`zip` where iterables must be exact.
+
+    ``strict`` option was added in 3.10. This function is always strict.
+    Args:
+        iterables (object): things to be zipped
+    Returns:
+        object: Iterator for zipping
+    """
+    if _version_ok(3, 10):
+        return zip(*iterables, strict=True)
+    return _zip_strict(iterables)
+
+
 def _assert_type(value, typ):
     if not isinstance(value, typ):
         raise TypeError(
@@ -99,3 +119,32 @@ def _assert_type(value, typ):
                 type(value),
             ),
         )
+
+
+def _version_ok(major, minor):
+    return sys.version_info >= (major, minor)
+
+
+def _zip_strict(iterables):
+    """Code from https://stackoverflow.com/a/69485272"""
+
+    def _first_tail():
+        nonlocal first_stopped
+        first_stopped = True
+        return
+        yield
+
+    def _zip_tail():
+        if not first_stopped:
+            raise ValueError("first iterable is longer than rest")
+        for x in itertools.chain.from_iterable(rest):
+            raise ValueError(f"left over data={str(x):.100} in iterable")
+            yield
+
+    if len(iterables) < 2:
+        return zip(*iterables)
+    first_stopped = False
+    iterables = iter(iterables)
+    first = itertools.chain(next(iterables), _first_tail())
+    rest = list(map(iter, iterables))
+    return itertools.chain(zip(first, *rest), _zip_tail())
diff --git a/pykern/pkjson.py b/pykern/pkjson.py
index a194e87..6f01b11 100644
--- a/pykern/pkjson.py
+++ b/pykern/pkjson.py
@@ -1,16 +1,22 @@
 # -*- coding: utf-8 -*-
 """JSON wrapper
 
-:copyright: Copyright (c) 2017 RadiaSoft LLC.  All Rights Reserved.
+:copyright: Copyright (c) 2017-2023 RadiaSoft LLC.  All Rights Reserved.
 :license: http://www.apache.org/licenses/LICENSE-2.0.html
 """
-from __future__ import absolute_import, division, print_function
 import json
 
 
 #: how bytes are encoded
 ENCODING = "utf-8"
 
+#: MIME type
+MIME_TYPE = "application/json"
+
+
+#: Content-Type MIME header
+CONTENT_TYPE = f'{MIME_TYPE}; charset="{ENCODING}"'
+
 
 class Encoder(json.JSONEncoder):
     def default(self, obj):
@@ -41,7 +47,7 @@ def dump_pretty(obj, filename=None, pretty=True, **kwargs):
                 separators=(",", ": "),
                 sort_keys=True,
                 cls=Encoder,
-                **kwargs
+                **kwargs,
             )
             + "\n"
         )


Test patch:
diff --git a/tests/pkcompat_test.py b/tests/pkcompat_test.py
index 211a840..9a51e2f 100644
--- a/tests/pkcompat_test.py
+++ b/tests/pkcompat_test.py
@@ -1,17 +1,12 @@
 # -*- coding: utf-8 -*-
-"""pytest for :mod:`pykern.pkcompat`
+"""test for :mod:`pykern.pkcompat`
 
 :copyright: Copyright (c) 2015 Bivio Software, Inc.  All Rights Reserved.
 :license: http://www.apache.org/licenses/LICENSE-2.0.html
 """
-from __future__ import absolute_import, division, print_function
-
 import locale
 import os
 
-import pytest
-import six
-
 
 def setup_module():
     """Set locale so can test expected outputs.
@@ -31,10 +26,7 @@ def test_from_bytes():
     s = pkcompat.from_bytes(b)
     pkeq(s, "你好")
     pkeq(b, b"\xe4\xbd\xa0\xe5\xa5\xbd")
-    if six.PY2:
-        pkeq(b, s)
-    else:
-        pkeq(False, b == s)
+    pkeq(False, b == s)
 
 
 def test_locale_str_1():
@@ -42,44 +34,23 @@ def test_locale_str_1():
     from pykern import pkcompat
 
     s = pkcompat.locale_str(b"\xc2\xb0")
-    if six.PY2:
-        assert isinstance(
-            s, unicode
-        ), "When locale_str is converted in PY2, it should return unicode"
-    else:
-        assert isinstance(
-            s, str
-        ), "When locale_str is converted in not PY2, it should return str"
+    assert isinstance(s, str), "When locale_str is converted, it should return str"
     assert "°" == s, "Conversion should be same as literal unicode value"
-    if six.PY2:
-        before = unicode(b"\xc2\xb0", "utf8")
-        assert before == pkcompat.locale_str(
-            before
-        ), "When string is already unicode, conversion yields same string"
-        before = str(123)
-        assert unicode(before) == pkcompat.locale_str(
-            before
-        ), "When string is already unicode, conversion yields same string"
-        before = str(None)
-        assert unicode(before) == pkcompat.locale_str(
-            before
-        ), "When string is already unicode, conversion yields same string"
-    else:
-        before = str(123)
-        assert before == pkcompat.locale_str(
-            before
-        ), "When string is already unicode, conversion yields same string"
-        before = str(None)
-        assert before == pkcompat.locale_str(
-            before
-        ), "When string is already unicode, conversion yields same string"
+    before = str(123)
+    assert before == pkcompat.locale_str(
+        before
+    ), "When string is already unicode, conversion yields same string"
+    before = str(None)
+    assert before == pkcompat.locale_str(
+        before
+    ), "When string is already unicode, conversion yields same string"
 
 
 def test_locale_str_2():
     """Invalid utf8"""
-    from pykern import pkcompat
+    from pykern import pkcompat, pkunit
 
-    with pytest.raises(UnicodeDecodeError):
+    with pkunit.pkexcept(UnicodeDecodeError):
         # TODO(robngler) set the locale?
         pkcompat.locale_str(b"\x80")
 
@@ -88,3 +59,26 @@ def test_unicode_unescape():
     from pykern import pkcompat
 
     assert "\n" == pkcompat.unicode_unescape(r"\n")
+
+
+def test_zip_strict():
+    from pykern import pkcompat, pkunit
+
+    for e, c in (
+        ((), ()),
+        ((), (range(0), range(0))),
+        (((0,),), (range(1),)),
+        (((0, 0),), (range(1), range(1))),
+        (((0, 0, 0),), (range(1), range(1), range(1))),
+    ):
+        pkunit.pkeq(e, tuple(pkcompat.zip_strict(*c)))
+
+    for c in (
+        (range(0), range(1)),
+        (range(0), range(0), range(1)),
+        (range(1), range(0)),
+        (range(1), range(2)),
+        (range(2), range(3)),
+    ):
+        with pkunit.pkexcept(ValueError):
+            list(pkcompat.zip_strict(*c))