refactor: clean up Plugin init ownership logic and format tests

nilslice · nilslice · commit 9a156f11a451 · 2026-02-08T18:18:29.000-08:00
Restructure Plugin.__init__ to make compiled_plugin ownership more
explicit with type annotation, and apply consistent formatting to tests.
diff --git a/extism/extism.py b/extism/extism.py
@@ -552,11 +552,13 @@ def __init__(
         functions: Optional[List[Function]] = HOST_FN_REGISTRY,
     ):
         # Track if we created the CompiledPlugin (so we know to free it)
-        self._owns_compiled_plugin = not isinstance(plugin, CompiledPlugin)
-        if self._owns_compiled_plugin:
-            plugin = CompiledPlugin(plugin, wasi, functions)
-
-        self.compiled_plugin = plugin
+        if isinstance(plugin, CompiledPlugin):
+            self._owns_compiled_plugin = False
+            self.compiled_plugin: Optional[CompiledPlugin] = plugin
+        else:
+            self._owns_compiled_plugin = True
+            self.compiled_plugin = CompiledPlugin(plugin, wasi, functions)
+        assert self.compiled_plugin is not None
         errmsg = _ffi.new("char**")
 
         self.plugin = _lib.extism_plugin_new_from_compiled(
@@ -632,7 +634,10 @@ def __del__(self):
         _lib.extism_plugin_free(self.plugin)
         self.plugin = -1
         # Free the compiled plugin if we created it
-        if getattr(self, "_owns_compiled_plugin", False) and self.compiled_plugin is not None:
+        if (
+            getattr(self, "_owns_compiled_plugin", False)
+            and self.compiled_plugin is not None
+        ):
             self.compiled_plugin.__del__()
             self.compiled_plugin = None
 
diff --git a/tests/test_extism.py b/tests/test_extism.py
@@ -52,11 +52,11 @@ def test_can_free_plugin(self):
 
     def test_plugin_del_frees_native_resources(self):
         """Test that Plugin.__del__ properly frees native resources.
-        
+
         This tests the fix for a bug where Plugin.__del__ checked for
         'self.pointer' instead of 'self.plugin', causing extism_plugin_free
         to never be called and leading to memory leaks.
-        
+
         This also tests that __del__ can be safely called multiple times
         (via context manager exit and garbage collection) without causing
         double-free errors.
@@ -66,66 +66,80 @@ def test_plugin_del_frees_native_resources(self):
             self.assertEqual(j["count"], 1)
             # Plugin should own the compiled plugin it created
             self.assertTrue(plugin._owns_compiled_plugin)
-        
+
         # Verify plugin was freed after exiting context
-        self.assertEqual(plugin.plugin, -1, 
-            "Expected plugin.plugin to be -1 after __del__, indicating extism_plugin_free was called")
+        self.assertEqual(
+            plugin.plugin,
+            -1,
+            "Expected plugin.plugin to be -1 after __del__, indicating extism_plugin_free was called",
+        )
         # Verify compiled plugin was also freed (since Plugin owned it)
-        self.assertIsNone(plugin.compiled_plugin,
-            "Expected compiled_plugin to be None after __del__, indicating it was also freed")
+        self.assertIsNone(
+            plugin.compiled_plugin,
+            "Expected compiled_plugin to be None after __del__, indicating it was also freed",
+        )
 
     def test_compiled_plugin_del_frees_native_resources(self):
         """Test that CompiledPlugin.__del__ properly frees native resources.
-        
+
         Unlike Plugin, CompiledPlugin has no context manager so __del__ is only
         called once by garbage collection. This also tests that __del__ can be
         safely called multiple times without causing double-free errors.
         """
         compiled = CompiledPlugin(self._manifest(), functions=[])
         # Verify pointer exists before deletion
-        self.assertTrue(hasattr(compiled, 'pointer'))
+        self.assertTrue(hasattr(compiled, "pointer"))
         self.assertNotEqual(compiled.pointer, -1)
-        
+
         # Create a plugin from compiled to ensure it works
         plugin = extism.Plugin(compiled)
         j = json.loads(plugin.call("count_vowels", "test"))
         self.assertEqual(j["count"], 1)
-        
+
         # Plugin should NOT own the compiled plugin (it was passed in)
         self.assertFalse(plugin._owns_compiled_plugin)
-        
+
         # Clean up plugin first
         plugin.__del__()
         self.assertEqual(plugin.plugin, -1)
-        
+
         # Compiled plugin should NOT have been freed by Plugin.__del__
-        self.assertNotEqual(compiled.pointer, -1,
-            "Expected compiled.pointer to NOT be -1 since Plugin didn't own it")
-        
+        self.assertNotEqual(
+            compiled.pointer,
+            -1,
+            "Expected compiled.pointer to NOT be -1 since Plugin didn't own it",
+        )
+
         # Now clean up compiled plugin manually
         compiled.__del__()
-        
+
         # Verify compiled plugin was freed
-        self.assertEqual(compiled.pointer, -1,
-            "Expected compiled.pointer to be -1 after __del__, indicating extism_compiled_plugin_free was called")
+        self.assertEqual(
+            compiled.pointer,
+            -1,
+            "Expected compiled.pointer to be -1 after __del__, indicating extism_compiled_plugin_free was called",
+        )
 
     def test_extism_function_metadata_del_frees_native_resources(self):
         """Test that _ExtismFunctionMetadata.__del__ properly frees native resources."""
+
         def test_host_fn(inp: str) -> str:
             return inp
-        
+
         func = TypeInferredFunction(None, "test_func", test_host_fn, [])
         metadata = _ExtismFunctionMetadata(func)
-        
+
         # Verify pointer exists before deletion
-        self.assertTrue(hasattr(metadata, 'pointer'))
+        self.assertTrue(hasattr(metadata, "pointer"))
         self.assertIsNotNone(metadata.pointer)
-        
+
         metadata.__del__()
-        
+
         # Verify function was freed (pointer set to None)
-        self.assertIsNone(metadata.pointer,
-            "Expected metadata.pointer to be None after __del__, indicating extism_function_free was called")
+        self.assertIsNone(
+            metadata.pointer,
+            "Expected metadata.pointer to be None after __del__, indicating extism_function_free was called",
+        )
 
     def test_errors_on_bad_manifest(self):
         self.assertRaises(
