From fc617142c02cfbe26711f310d3f52bc9703c9eb8 Mon Sep 17 00:00:00 2001
From: Madhava Jay <me@madhavajay.com>
Date: Tue, 24 May 2022 11:38:27 +1000
Subject: [PATCH] Fixed issue where pickle helpers didnt use new from_bytes
 with context

- Added with context to benchmarks that use with_bytes
---
 benchmark/addressbook.capnp.orphan.py | 13 ++++++-------
 benchmark/addressbook.capnp.py        | 11 +++++------
 capnp/lib/capnp.pyx                   |  3 ++-
 capnp/lib/pickle_helper.py            |  3 ++-
 scripts/capnp_test_pycapnp.py         |  3 ++-
 5 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/benchmark/addressbook.capnp.orphan.py b/benchmark/addressbook.capnp.orphan.py
index 4717884..4ac6e59 100644
--- a/benchmark/addressbook.capnp.orphan.py
+++ b/benchmark/addressbook.capnp.orphan.py
@@ -35,13 +35,12 @@ def writeAddressBook():
 
 
 def printAddressBook(msg_bytes):
-    addressBook = addressbook.AddressBook.from_bytes(msg_bytes)
-
-    for person in addressBook.people:
-        print(person.name, ":", person.email)
-        for phone in person.phones:
-            print(phone.type, ":", phone.number)
-        print()
+    with addressbook.AddressBook.from_bytes(msg_bytes) as addressBook:
+        for person in addressBook.people:
+            print(person.name, ":", person.email)
+            for phone in person.phones:
+                print(phone.type, ":", phone.number)
+            print()
 
 
 if __name__ == "__main__":
diff --git a/benchmark/addressbook.capnp.py b/benchmark/addressbook.capnp.py
index 04c3370..83c6ba1 100644
--- a/benchmark/addressbook.capnp.py
+++ b/benchmark/addressbook.capnp.py
@@ -40,12 +40,11 @@ def writeAddressBook():
 
 @profile
 def printAddressBook(msg_bytes):
-    addressBook = addressbook.AddressBook.from_bytes(msg_bytes)
-
-    for person in addressBook.people:
-        person.name, person.email
-        for phone in person.phones:
-            phone.type, phone.number
+    with addressbook.AddressBook.from_bytes(msg_bytes) as addressBook:
+        for person in addressBook.people:
+            person.name, person.email
+            for phone in person.phones:
+                phone.type, phone.number
 
 
 @profile
diff --git a/capnp/lib/capnp.pyx b/capnp/lib/capnp.pyx
index e277ad9..19995a6 100644
--- a/capnp/lib/capnp.pyx
+++ b/capnp/lib/capnp.pyx
@@ -1099,7 +1099,8 @@ if getattr(_sys, 'subversion', [''])[0] == 'PyPy':
     from pickle_helper import _struct_reducer
 else:
     def _struct_reducer(schema_id, data):
-        return _global_schema_parser.modules_by_id[schema_id].from_bytes(data)
+        with _global_schema_parser.modules_by_id[schema_id].from_bytes(data) as msg:
+            return msg
 
 
 cdef class _DynamicStructReader:
diff --git a/capnp/lib/pickle_helper.py b/capnp/lib/pickle_helper.py
index c28d46a..fbf6a91 100644
--- a/capnp/lib/pickle_helper.py
+++ b/capnp/lib/pickle_helper.py
@@ -3,4 +3,5 @@ import capnp
 
 def _struct_reducer(schema_id, data):
     'Hack to deal with pypy not allowing reduce functions to be "built-in" methods (ie. compiled from a .pyx)'
-    return capnp._global_schema_parser.modules_by_id[schema_id].from_bytes(data)
+    with capnp._global_schema_parser.modules_by_id[schema_id].from_bytes(data) as msg:
+        return msg
diff --git a/scripts/capnp_test_pycapnp.py b/scripts/capnp_test_pycapnp.py
index 531aa40..897ee36 100755
--- a/scripts/capnp_test_pycapnp.py
+++ b/scripts/capnp_test_pycapnp.py
@@ -13,7 +13,8 @@ import test_capnp  # noqa: E402
 
 def decode(name):
     class_name = name[0].upper() + name[1:]
-    print(getattr(test_capnp, class_name).from_bytes(sys.stdin.read())._short_str())
+    with getattr(test_capnp, class_name).from_bytes(sys.stdin.read()) as msg:
+        print(msg._short_str())
 
 
 def encode(name):