Răsfoiți Sursa

Add resource management to Bundle class

Jelmer Vernooij 2 săptămâni în urmă
părinte
comite
b9e61f621b
3 a modificat fișierele cu 46 adăugiri și 0 ștergeri
  1. 29 0
      dulwich/bundle.py
  2. 7 0
      tests/compat/test_bundle.py
  3. 10 0
      tests/test_bundle.py

+ 29 - 0
dulwich/bundle.py

@@ -101,6 +101,31 @@ class Bundle:
             return False
         return True
 
+    def close(self) -> None:
+        """Close any open resources in this bundle."""
+        if self.pack_data is not None:
+            self.pack_data.close()
+            self.pack_data = None
+
+    def __enter__(self) -> "Bundle":
+        """Enter context manager."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit context manager and close bundle."""
+        self.close()
+
+    def __del__(self) -> None:
+        """Warn if bundle was not explicitly closed."""
+        if self.pack_data is not None:
+            import warnings
+            warnings.warn(
+                f"Bundle {self!r} was not explicitly closed. "
+                "Please use bundle.close() or a context manager.",
+                ResourceWarning,
+                stacklevel=2,
+            )
+
     def store_objects(
         self,
         object_store: "BaseObjectStore",
@@ -332,6 +357,10 @@ def create_bundle_from_repo(
         def iter_unpacked(self) -> Iterator[UnpackedObject]:
             return iter(self._objects)
 
+        def close(self) -> None:
+            """Close pack data (no-op for in-memory pack data)."""
+            pass
+
     pack_data = _BundlePackData(pack_count, pack_objects, repo.object_format)
 
     # Create bundle object

+ 7 - 0
tests/compat/test_bundle.py

@@ -62,6 +62,7 @@ class CompatBundleTestCase(CompatTestCase):
 
         # Use create_bundle_from_repo helper
         bundle = create_bundle_from_repo(self.repo)
+        self.addCleanup(bundle.close)
 
         with open(bundle_path, "wb") as f:
             write_bundle(f, bundle)
@@ -94,6 +95,7 @@ class CompatBundleTestCase(CompatTestCase):
         # Read bundle using dulwich
         with open(bundle_path, "rb") as f:
             bundle = read_bundle(f)
+        self.addCleanup(bundle.close)
 
         # Verify bundle contents
         self.assertEqual(2, bundle.version)
@@ -138,6 +140,7 @@ class CompatBundleTestCase(CompatTestCase):
         # Read bundle using dulwich
         with open(bundle_path, "rb") as f:
             bundle = read_bundle(f)
+        self.addCleanup(bundle.close)
 
         # Verify bundle contains all refs
         self.assertIn(b"refs/heads/master", bundle.references)
@@ -185,6 +188,7 @@ class CompatBundleTestCase(CompatTestCase):
         # Read bundle using dulwich
         with open(bundle_path, "rb") as f:
             bundle = read_bundle(f)
+        self.addCleanup(bundle.close)
 
         # Verify bundle has prerequisites
         self.assertGreater(len(bundle.prerequisites), 0)
@@ -234,6 +238,7 @@ class CompatBundleTestCase(CompatTestCase):
         # Read bundle using dulwich
         with open(bundle_path, "rb") as f:
             bundle = read_bundle(f)
+        self.addCleanup(bundle.close)
 
         # Verify bundle contents
         self.assertEqual(2, bundle.version)
@@ -266,6 +271,7 @@ class CompatBundleTestCase(CompatTestCase):
         # Read bundle using dulwich
         with open(bundle_path, "rb") as f:
             bundle = read_bundle(f)
+        self.addCleanup(bundle.close)
 
         # Verify bundle was read correctly
         self.assertEqual(2, bundle.version)
@@ -319,6 +325,7 @@ class CompatBundleTestCase(CompatTestCase):
         # Read the bundle and store objects using dulwich
         with open(bundle_path, "rb") as f:
             bundle = read_bundle(f)
+        self.addCleanup(bundle.close)
 
         # Use the bundle's store_objects method to unbundle
         bundle.store_objects(unbundle_repo.object_store)

+ 10 - 0
tests/test_bundle.py

@@ -53,6 +53,7 @@ class BundleTests(TestCase):
         write_pack_objects(b.write, [], object_format=DEFAULT_OBJECT_FORMAT)
         b.seek(0)
         bundle.pack_data = PackData.from_file(b, object_format=DEFAULT_OBJECT_FORMAT)
+        self.addCleanup(bundle.pack_data.close)
 
         # Check the repr output
         rep = repr(bundle)
@@ -159,6 +160,7 @@ class BundleTests(TestCase):
         f.seek(0)
 
         bundle = read_bundle(f)
+        self.addCleanup(bundle.close)
         self.assertEqual(2, bundle.version)
         self.assertEqual({}, bundle.capabilities)
         self.assertEqual([(b"cc" * 20, b"prerequisite comment")], bundle.prerequisites)
@@ -307,6 +309,7 @@ class BundleTests(TestCase):
 
     def test_roundtrip_bundle(self) -> None:
         origbundle = Bundle()
+        self.addCleanup(origbundle.close)
         origbundle.version = 3
         origbundle.capabilities = {"foo": None}
         origbundle.references = {b"refs/heads/master": b"ab" * 20}
@@ -323,6 +326,7 @@ class BundleTests(TestCase):
 
             with open(os.path.join(td, "foo"), "rb") as f:
                 newbundle = read_bundle(f)
+                self.addCleanup(newbundle.close)
 
                 self.assertEqual(origbundle, newbundle)
 
@@ -355,6 +359,7 @@ class BundleTests(TestCase):
 
         # Create bundle from repository
         bundle = create_bundle_from_repo(repo)
+        self.addCleanup(bundle.close)
 
         # Verify bundle contents
         self.assertEqual(bundle.references, {b"refs/heads/master": commit.id})
@@ -394,6 +399,7 @@ class BundleTests(TestCase):
         # Create bundle with prerequisites
         prereq_id = b"aa" * 20  # hex string like other object ids
         bundle = create_bundle_from_repo(repo, prerequisites=[prereq_id])
+        self.addCleanup(bundle.close)
 
         # Verify prerequisites are included
         self.assertEqual(len(bundle.prerequisites), 1)
@@ -426,6 +432,7 @@ class BundleTests(TestCase):
         from dulwich.refs import Ref
 
         bundle = create_bundle_from_repo(repo, refs=[Ref(b"refs/heads/master")])
+        self.addCleanup(bundle.close)
 
         # Verify only master ref is included
         self.assertEqual(len(bundle.references), 1)
@@ -457,6 +464,7 @@ class BundleTests(TestCase):
         # Create bundle with capabilities
         capabilities = {"object-format": "sha1"}
         bundle = create_bundle_from_repo(repo, capabilities=capabilities, version=3)
+        self.addCleanup(bundle.close)
 
         # Verify capabilities are included
         self.assertEqual(bundle.capabilities, capabilities)
@@ -489,6 +497,7 @@ class BundleTests(TestCase):
 
         # Use blob.id directly (40-byte hex bytestring)
         bundle = create_bundle_from_repo(repo, prerequisites=[prereq_blob.id])
+        self.addCleanup(bundle.close)
 
         # Verify the prerequisite was added correctly
         self.assertEqual(len(bundle.prerequisites), 1)
@@ -520,6 +529,7 @@ class BundleTests(TestCase):
         prereq_hex = b"aa" * 20
 
         bundle = create_bundle_from_repo(repo, prerequisites=[prereq_hex])
+        self.addCleanup(bundle.close)
 
         # Verify the prerequisite was added correctly
         self.assertEqual(len(bundle.prerequisites), 1)