pythonnet · lostmsu · May 10, 2024 · Feb 6, 2024 · Feb 21, 2024 · Mar 17, 2024
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,6 +10,9 @@ This document follows the conventions laid out in [Keep a CHANGELOG][].
 ### Added
 
 ### Changed
+-   Added a `FormatterFactory` member in RuntimeData to create formatters with parameters. For compatibility, the `FormatterType` member is still present and has precedence when defining both `FormatterFactory` and `FormatterType`
+-   Added a post-serialization and a pre-deserialization step callbacks to extend (de)serialization process
+-   Added an API to stash serialized data on Python capsules
 
 ### Fixed
 

diff --git a/src/runtime/StateSerialization/RuntimeData.cs b/src/runtime/StateSerialization/RuntimeData.cs
@@ -17,7 +17,27 @@ namespace Python.Runtime
 {
     public static class RuntimeData
     {
-        private static Type? _formatterType;
+
+        private readonly static Func<IFormatter> DefaultFormatter = () => new BinaryFormatter();
+        private static Func<IFormatter>? _formatterFactory { get; set; } = null;
+
+        public static Func<IFormatter> FormatterFactory
+        {
+            get
+            {
+                if (_formatterFactory is null)
+                {
+                    return DefaultFormatter;
+                }
+                return _formatterFactory;
+            }
+            set
+            {
+                _formatterFactory = value;
+            }
+        }
+
+        private static Type? _formatterType = null;
         public static Type? FormatterType
         {
             get => _formatterType;
@@ -31,6 +51,14 @@ public static Type? FormatterType
             }
         }
 
+        /// <summary>
+        /// Callback called as a last step in the serialization process
+        /// </summary>
+        public static Action? PostStashHook {get; set;} = null;
+        /// <summary>
+        /// Callback called as the first step in the deserialization process
+        /// </summary>
+        public static Action? PreRestoreHook {get; set;} = null;
         public static ICLRObjectStorer? WrappersStorer { get; set; }
 
         /// <summary>
@@ -74,6 +102,7 @@ internal static void Stash()
             using NewReference capsule = PyCapsule_New(mem, IntPtr.Zero, IntPtr.Zero);
             int res = PySys_SetObject("clr_data", capsule.BorrowOrThrow());
             PythonException.ThrowIfIsNotZero(res);
+            PostStashHook?.Invoke();
         }
 
         internal static void RestoreRuntimeData()
@@ -90,6 +119,7 @@ internal static void RestoreRuntimeData()
 
         private static void RestoreRuntimeDataImpl()
         {
+            PreRestoreHook?.Invoke();
             BorrowedReference capsule = PySys_GetObject("clr_data");
             if (capsule.IsNull)
             {
@@ -250,11 +280,82 @@ private static void RestoreRuntimeDataObjects(SharedObjectsState storage)
             }
         }
 
+        /// <summary>
+        /// Frees the pointer stored in a Python capsule and the capsule stored 
+        /// on the sys module object with the given key if it exists. 
+        /// </summary>
+        /// <remarks>
+        /// The memory on the capsule must have been allocated via <code>StashDataInCapsule</code>
+        /// </remarks>
+        /// <param name="key">The name given to the capsule on the `sys` module object</param>
+        public static void FreeCapsuleData(string key)
+        {
+            BorrowedReference oldCapsule = PySys_GetObject(key);
+            if (!oldCapsule.IsNull)
+            {
+                IntPtr oldData = PyCapsule_GetPointer(oldCapsule, IntPtr.Zero);
+                Marshal.FreeHGlobal(oldData);
+                PyCapsule_SetPointer(oldCapsule, IntPtr.Zero);
+                PySys_SetObject(key, null);
+            }
+        }
+
+        /// <summary>
+        /// Stores the <paramref name="data"/>data parameter in a Python capsule and stores 
+        /// the capsule on the `sys` module object with the name <paramref name="key"/>. 
+        /// This method allocates global memory to hold the data of the <paramref name="data"/>
+        /// parameter.
+        /// </summary>
+        /// <remarks>
+        /// No checks on pre-existing names on the `sys` module object are made.
+        /// </remarks>
+        /// <param name="key">The name given to the capsule on the `sys` module object</param>
+        /// <param name="data">The data to be contained in the capsule</param>
+        public static void StashDataInCapsule(string key, byte[] data)
+        {
+            IntPtr mem = Marshal.AllocHGlobal(IntPtr.Size + data.Length);
+            // store the length of the buffer first
+            Marshal.WriteIntPtr(mem, (IntPtr)data.Length);
+            Marshal.Copy(data, 0, mem + IntPtr.Size, data.Length);
+
+            using NewReference capsule = PyCapsule_New(mem, IntPtr.Zero, IntPtr.Zero);
+            int res = PySys_SetObject(key, capsule.BorrowOrThrow());
+            PythonException.ThrowIfIsNotZero(res);
+        }
+
+        /// <summary>
+        /// Retreives the pointer to previously stored data on a Python capsule. 
+        /// Throws if the object corresponding to the <paramref name="key"/> parameter
+        /// on the `sys` module object is not a capsule.
+        /// </summary>
+        /// <param name="key">The name given to the capsule on the `sys` module object</param>
+        /// <returns>The pointer to the data, or IntPtr.Zero if name matches the key</returns>
+        public static IntPtr GetDataFromCapsule(string key)
+        {
+            BorrowedReference capsule = PySys_GetObject(key);
+            if (capsule.IsNull)
+            {
+                // nothing to do.
+                return IntPtr.Zero;
+            }
+            var ptr = PyCapsule_GetPointer(capsule, IntPtr.Zero);
+            if (ptr == IntPtr.Zero)
+            {
+                // The PyCapsule API returns NULL on error; NULL cannot be stored 
+                // as a capsule's value
+                PythonException.ThrowIfIsNull(null);
+            }
+            return ptr;
+        }
+
         internal static IFormatter CreateFormatter()
         {
-            return FormatterType != null ?
-                (IFormatter)Activator.CreateInstance(FormatterType)
-                : new BinaryFormatter();
+
+            if (FormatterType != null)
+            {
+                return (IFormatter)Activator.CreateInstance(FormatterType);
+            }
+            return FormatterFactory();
         }
     }
 }