From 57017b6526bdbf81c52b67e79a8de3acd326070c Mon Sep 17 00:00:00 2001
From: Alexey Khudyakov <alexey.skladnoy@gmail.com>
Date: Thu, 13 Feb 2025 00:08:23 +0300
Subject: [PATCH] Add ToPy instances for Py and functions returning Py

---
 src/Python/Inline/Literal.hs | 35 ++++++++++++-
 test/TST/Callbacks.hs        | 97 ++++++++++++++++++++----------------
 2 files changed, 88 insertions(+), 44 deletions(-)

diff --git a/src/Python/Inline/Literal.hs b/src/Python/Inline/Literal.hs
index 72fbbbd..959dd1b 100644
--- a/src/Python/Inline/Literal.hs
+++ b/src/Python/Inline/Literal.hs
@@ -577,7 +577,6 @@ instance (ToPy b) => ToPy (IO b) where
     --
     [CU.exp| PyObject* { inline_py_callback_METH_NOARGS($(PyCFunction f_ptr)) } |]
 
-
 -- | Only accepts positional parameters
 instance (FromPy a, Show a, ToPy b) => ToPy (a -> IO b) where
   basicToPy f = Py $ do
@@ -600,6 +599,40 @@ instance (FromPy a1, FromPy a2, ToPy b) => ToPy (a1 -> a2 -> IO b) where
     --
     [CU.exp| PyObject* { inline_py_callback_METH_FASTCALL($(PyCFunctionFast f_ptr)) } |]
 
+
+-- | Converted to 0-ary function
+instance (ToPy b) => ToPy (Py b) where
+  basicToPy f = Py $ do
+    --
+    f_ptr <- wrapCFunction $ \_ _ -> pyCallback $ do
+      progPy $ basicToPy =<< f
+    --
+    [CU.exp| PyObject* { inline_py_callback_METH_NOARGS($(PyCFunction f_ptr)) } |]
+
+-- | Only accepts positional parameters
+instance (FromPy a, Show a, ToPy b) => ToPy (a -> Py b) where
+  basicToPy f = Py $ do
+    --
+    f_ptr <- wrapCFunction $ \_ p_a -> pyCallback $ do
+      a <- loadArg p_a 0 1
+      progPy $ basicToPy =<< f a
+    --
+    [CU.exp| PyObject* { inline_py_callback_METH_O($(PyCFunction f_ptr)) } |]
+
+-- | Only accepts positional parameters
+instance (FromPy a1, FromPy a2, ToPy b) => ToPy (a1 -> a2 -> Py b) where
+  basicToPy f = Py $ do
+    --
+    f_ptr <- wrapFastcall $ \_ p_arr n -> pyCallback $ do
+      when (n /= 2) $ abortM $ raiseBadNArgs 2 n
+      a1 <- loadArgFastcall p_arr 0 n
+      a2 <- loadArgFastcall p_arr 1 n
+      progPy $ basicToPy =<< f a1 a2
+    --
+    [CU.exp| PyObject* { inline_py_callback_METH_FASTCALL($(PyCFunctionFast f_ptr)) } |]
+
+
+
 ----------------------------------------------------------------
 -- Helpers
 ----------------------------------------------------------------
diff --git a/test/TST/Callbacks.hs b/test/TST/Callbacks.hs
index b09e7b5..adc4dae 100644
--- a/test/TST/Callbacks.hs
+++ b/test/TST/Callbacks.hs
@@ -54,48 +54,59 @@ tests = testGroup "Callbacks"
          except TypeError as e:
              pass
          |]
-  , testCase "Haskell exception in callback(arity=1)" $ runPy $ do
-      let foo :: Int -> IO Int
-          foo y = pure $ 10 `div` y
-      throwsPy [py_| foo_hs(0) |]
-  , testCase "Haskell exception in callback(arity=2)" $ runPy $ do
-      let foo :: Int -> Int -> IO Int
-          foo x y = pure $ x `div` y
-      throwsPy [py_| foo_hs(1, 0) |]
     ----------------------------------------
-  , testCase "Call python in callback (arity=1)" $ runPy $ do
-      let foo :: Int -> IO Int
-          foo x = do Just x' <- runPy $ fromPy =<< [pye| 100 // x_hs |]
-                     pure x'
-      [py_|
-        assert foo_hs(5) == 20
-        |]
-  , testCase "Call python in callback (arity=2" $ runPy $ do
-      let foo :: Int -> Int -> IO Int
-          foo x y = do Just x' <- runPy $ fromPy =<< [pye| x_hs // y_hs |]
-                       pure x'
-      [py_|
-        assert foo_hs(100,5) == 20
-        |]
+  , testCase "Py function(arity 0)" $ do
+      let fun = [pye| 0 |]
+      runPy [py_| assert fun_hs() == 0 |]
+  , testCase "Py function(arity=1)" $ runPy $ do
+      let double (n::Int) = [pye| n_hs * 2 |]
+      [py_| assert double_hs(3) == 6 |]
+  , testCase "Py function(arity=2)" $ runPy $ do
+      let foo (x::Int) (y::Int) = [pye| x_hs * y_hs |]
+      [py_| assert foo_hs(3, 100) == 300 |]
     ----------------------------------------
-  , testCase "No leaks (arity=1)" $ runPy $ do
-      let foo :: Int -> IO Int
-          foo y = pure $ 10 * y
-      [py_|
-        import sys
-        x = 123456
-        old_refcount = sys.getrefcount(x)
-        foo_hs(x)
-        assert old_refcount == sys.getrefcount(x)
-        |]
-  , testCase "No leaks (arity=2)" $ runPy $ do
-      let foo :: Int -> Int -> IO Int
-          foo x y = pure $ x * y
-      [py_|
-        import sys
-        x = 123456
-        old_refcount = sys.getrefcount(x)
-        foo_hs(1,x)
-        assert old_refcount == sys.getrefcount(x)
-        |]
-  ]
+  , testCase "Haskell exception in callback(arity=1)" $ runPy $ do
+       let foo :: Int -> IO Int
+           foo y = pure $ 10 `div` y
+       throwsPy [py_| foo_hs(0) |]
+   , testCase "Haskell exception in callback(arity=2)" $ runPy $ do
+       let foo :: Int -> Int -> IO Int
+           foo x y = pure $ x `div` y
+       throwsPy [py_| foo_hs(1, 0) |]
+     ----------------------------------------
+   , testCase "Call python in callback (arity=1)" $ runPy $ do
+       let foo :: Int -> IO Int
+           foo x = do Just x' <- runPy $ fromPy =<< [pye| 100 // x_hs |]
+                      pure x'
+       [py_|
+         assert foo_hs(5) == 20
+         |]
+   , testCase "Call python in callback (arity=2" $ runPy $ do
+       let foo :: Int -> Int -> IO Int
+           foo x y = do Just x' <- runPy $ fromPy =<< [pye| x_hs // y_hs |]
+                        pure x'
+       [py_|
+         assert foo_hs(100,5) == 20
+         |]
+     ----------------------------------------
+   , testCase "No leaks (arity=1)" $ runPy $ do
+       let foo :: Int -> IO Int
+           foo y = pure $ 10 * y
+       [py_|
+         import sys
+         x = 123456
+         old_refcount = sys.getrefcount(x)
+         foo_hs(x)
+         assert old_refcount == sys.getrefcount(x)
+         |]
+   , testCase "No leaks (arity=2)" $ runPy $ do
+       let foo :: Int -> Int -> IO Int
+           foo x y = pure $ x * y
+       [py_|
+         import sys
+         x = 123456
+         old_refcount = sys.getrefcount(x)
+         foo_hs(1,x)
+         assert old_refcount == sys.getrefcount(x)
+         |]
+   ]