Fix hasndling of masking in Py moand and add instances for Set

inline-python.cabal

@@ -71,6 +71,7 @@ Library
     Python.Inline
     Python.Types
   Other-modules:
+    Python.Internal.CAPI
     Python.Internal.Eval
     Python.Internal.EvalQQ
     Python.Internal.Program
@@ -88,6 +89,7 @@ library test
                   , tasty-hunit      >=0.10
                   , tasty-quickcheck >=0.10
                   , exceptions
+                  , containers
   hs-source-dirs:   test
   Exposed-modules:
     TST.Run

src/Python/Inline/Literal.hs

@@ -14,13 +14,12 @@ module Python.Inline.Literal
 
 import Control.Monad
 import Control.Monad.Catch
-import Control.Monad.IO.Class
-import Control.Monad.Trans.Class
 import Control.Monad.Trans.Cont
 import Data.Bits
 import Data.Char
 import Data.Int
 import Data.Word
+import Data.Set                  qualified as Set
 import Foreign.Ptr
 import Foreign.C.Types
 import Foreign.Storable
@@ -32,7 +31,7 @@ import Language.C.Inline.Unsafe  qualified as CU
 import Python.Types
 import Python.Internal.Types
 import Python.Internal.Eval
-
+import Python.Internal.CAPI
 import Python.Internal.Program
 
 ----------------------------------------------------------------
@@ -56,17 +55,17 @@ class ToPy a where
   basicToPy :: a -> Py (Ptr PyObject)
   -- | Old hack for handling of strings
   basicListToPy :: [a] -> Py (Ptr PyObject)
-  basicListToPy xs = evalContT $ do
+  basicListToPy xs = runProgram $ do
     let n = fromIntegral $ length xs :: CLLong
-    p_list <- checkNull (Py [CU.exp| PyObject* { PyList_New($(long long n)) } |])
-    onExceptionProg $ decref p_list
-    let loop !_ []     = pure p_list
+    p_list <- takeOwnership =<< checkNull (Py [CU.exp| PyObject* { PyList_New($(long long n)) } |])
+    let loop !_ []     = p_list <$ incref p_list
         loop  i (a:as) = basicToPy a >>= \case
           NULL -> pure nullPtr
           p_a  -> do
-            liftIO [CU.exp| void { PyList_SET_ITEM($(PyObject* p_list), $(long long i), $(PyObject* p_a)) } |]
+            -- NOTE: PyList_SET_ITEM steals reference
+            Py [CU.exp| void { PyList_SET_ITEM($(PyObject* p_list), $(long long i), $(PyObject* p_a)) } |]
             loop (i+1) as
-    lift $ loop 0 xs
+    progPy $ loop 0 xs
 
 -- | Convert python object to haskell value.
 class FromPy a where
@@ -250,9 +249,9 @@ instance ToPy Char where
        uint32_t cs[1] = { $(unsigned i) };
        return PyUnicode_DecodeUTF32((char*)cs, 4, NULL, NULL);
        } |]
-  basicListToPy str = evalContT $ do
+  basicListToPy str = runProgram $ do
     p_str <- withPyWCString str
-    liftIO [CU.exp| PyObject* { PyUnicode_FromWideChar($(wchar_t *p_str), -1) } |]
+    progIO [CU.exp| PyObject* { PyUnicode_FromWideChar($(wchar_t *p_str), -1) } |]
 
 
 instance FromPy Char where
@@ -289,84 +288,84 @@ instance FromPy Bool where
 
 
 instance (ToPy a, ToPy b) => ToPy (a,b) where
-  basicToPy (a,b) = evalContT $ do
+  basicToPy (a,b) = runProgram $ do
     p_a <- takeOwnership =<< checkNull (basicToPy a)
     p_b <- takeOwnership =<< checkNull (basicToPy b)
-    liftIO [CU.exp| PyObject* { PyTuple_Pack(2, $(PyObject* p_a), $(PyObject* p_b)) } |]
+    progIO [CU.exp| PyObject* { PyTuple_Pack(2, $(PyObject* p_a), $(PyObject* p_b)) } |]
 
 -- | Will accept any iterable
 instance (FromPy a, FromPy b) => FromPy (a,b) where
-  basicFromPy p_tup = evalContT $ do
+  basicFromPy p_tup = runProgram $ do
     -- Unpack 2-tuple.
     p_args    <- withPyAllocaArray 2
-    unpack_ok <- liftIO [CU.exp| int {
+    unpack_ok <- progIO [CU.exp| int {
       inline_py_unpack_iterable($(PyObject *p_tup), 2, $(PyObject **p_args))
       }|]
-    lift $ do checkThrowPyError
-              when (unpack_ok /= 0) $ throwM BadPyType
+    progPy $ do checkThrowPyError
+                when (unpack_ok /= 0) $ throwM BadPyType
     -- Parse each element of tuple
-    p_a <- takeOwnership =<< liftIO (peekElemOff p_args 0)
-    p_b <- takeOwnership =<< liftIO (peekElemOff p_args 1)
-    lift $ do a <- basicFromPy p_a
-              b <- basicFromPy p_b
-              pure (a,b)
+    p_a <- takeOwnership =<< progIO (peekElemOff p_args 0)
+    p_b <- takeOwnership =<< progIO (peekElemOff p_args 1)
+    progPy $ do a <- basicFromPy p_a
+                b <- basicFromPy p_b
+                pure (a,b)
 
 instance (ToPy a, ToPy b, ToPy c) => ToPy (a,b,c) where
-  basicToPy (a,b,c) = evalContT $ do
+  basicToPy (a,b,c) = runProgram $ do
     p_a <- takeOwnership =<< checkNull (basicToPy a)
     p_b <- takeOwnership =<< checkNull (basicToPy b)
     p_c <- takeOwnership =<< checkNull (basicToPy c)
-    liftIO [CU.exp| PyObject* {
+    progIO [CU.exp| PyObject* {
       PyTuple_Pack(3, $(PyObject *p_a), $(PyObject *p_b), $(PyObject *p_c)) } |]
 
 -- | Will accept any iterable
 instance (FromPy a, FromPy b, FromPy c) => FromPy (a,b,c) where
-  basicFromPy p_tup = evalContT $ do
+  basicFromPy p_tup = runProgram $ do
     -- Unpack 3-tuple.
     p_args    <- withPyAllocaArray 3
-    unpack_ok <- liftIO [CU.exp| int {
+    unpack_ok <- progIO [CU.exp| int {
       inline_py_unpack_iterable($(PyObject *p_tup), 3, $(PyObject **p_args))
       }|]
-    lift $ do checkThrowPyError
-              when (unpack_ok /= 0) $ throwM BadPyType
+    progPy $ do checkThrowPyError
+                when (unpack_ok /= 0) $ throwM BadPyType
     -- Parse each element of tuple
-    p_a <- takeOwnership =<< liftIO (peekElemOff p_args 0)
-    p_b <- takeOwnership =<< liftIO (peekElemOff p_args 1)
-    p_c <- takeOwnership =<< liftIO (peekElemOff p_args 2)
-    lift $ do a <- basicFromPy p_a
-              b <- basicFromPy p_b
-              c <- basicFromPy p_c
-              pure (a,b,c)
+    p_a <- takeOwnership =<< progIO (peekElemOff p_args 0)
+    p_b <- takeOwnership =<< progIO (peekElemOff p_args 1)
+    p_c <- takeOwnership =<< progIO (peekElemOff p_args 2)
+    progPy $ do a <- basicFromPy p_a
+                b <- basicFromPy p_b
+                c <- basicFromPy p_c
+                pure (a,b,c)
 
 instance (ToPy a, ToPy b, ToPy c, ToPy d) => ToPy (a,b,c,d) where
-  basicToPy (a,b,c,d) = evalContT $ do
+  basicToPy (a,b,c,d) = runProgram $ do
     p_a <- takeOwnership =<< checkNull (basicToPy a)
     p_b <- takeOwnership =<< checkNull (basicToPy b)
     p_c <- takeOwnership =<< checkNull (basicToPy c)
     p_d <- takeOwnership =<< checkNull (basicToPy d)
-    liftIO [CU.exp| PyObject* {
+    progIO [CU.exp| PyObject* {
       PyTuple_Pack(4, $(PyObject *p_a), $(PyObject *p_b), $(PyObject *p_c), $(PyObject *p_d)) } |]
 
 -- | Will accept any iterable
 instance (FromPy a, FromPy b, FromPy c, FromPy d) => FromPy (a,b,c,d) where
-  basicFromPy p_tup = evalContT $ do
+  basicFromPy p_tup = runProgram $ do
     -- Unpack 3-tuple.
     p_args    <- withPyAllocaArray 4
-    unpack_ok <- liftIO [CU.exp| int {
+    unpack_ok <- progIO [CU.exp| int {
       inline_py_unpack_iterable($(PyObject *p_tup), 4, $(PyObject **p_args))
       }|]
-    lift $ do checkThrowPyError
-              when (unpack_ok /= 0) $ throwM BadPyType
+    progPy $ do checkThrowPyError
+                when (unpack_ok /= 0) $ throwM BadPyType
     -- Parse each element of tuple
-    p_a <- takeOwnership =<< liftIO (peekElemOff p_args 0)
-    p_b <- takeOwnership =<< liftIO (peekElemOff p_args 1)
-    p_c <- takeOwnership =<< liftIO (peekElemOff p_args 2)
-    p_d <- takeOwnership =<< liftIO (peekElemOff p_args 3)
-    lift $ do a <- basicFromPy p_a
-              b <- basicFromPy p_b
-              c <- basicFromPy p_c
-              d <- basicFromPy p_d
-              pure (a,b,c,d)
+    p_a <- takeOwnership =<< progIO (peekElemOff p_args 0)
+    p_b <- takeOwnership =<< progIO (peekElemOff p_args 1)
+    p_c <- takeOwnership =<< progIO (peekElemOff p_args 2)
+    p_d <- takeOwnership =<< progIO (peekElemOff p_args 3)
+    progPy $ do a <- basicFromPy p_a
+                b <- basicFromPy p_b
+                c <- basicFromPy p_c
+                d <- basicFromPy p_d
+                pure (a,b,c,d)
 
 instance (ToPy a) => ToPy [a] where
   basicToPy = basicListToPy
@@ -383,14 +382,45 @@ instance (FromPy a) => FromPy [a] where
       } |]
     when (nullPtr == p_iter) $ throwM BadPyType
     --
-    let loop f = do
-          p <- Py [C.exp| PyObject* { PyIter_Next($(PyObject* p_iter)) } |]
-          checkThrowPyError
-          case p of
-            NULL -> pure f
-            _    -> do a <- basicFromPy p `finally` decref p
-                       loop (f . (a:))
-    ($ []) <$> loop id
+    f <- foldPyIterable p_iter
+      (\f p -> do a <- basicFromPy p
+                  pure (f . (a:)))
+      id
+    pure $ f []
+
+instance (ToPy a, Ord a) => ToPy (Set.Set a) where
+  basicToPy set = runProgram $ do
+    p_set <- takeOwnership =<< checkNull basicNewSet
+    progPy $ do
+      let loop []     = p_set <$ incref p_set
+          loop (x:xs) = basicToPy x >>= \case
+            NULL -> pure NULL
+            p_a  -> Py [C.exp| int { PySet_Add($(PyObject *p_set), $(PyObject *p_a)) }|] >>= \case
+              0 -> decref p_a >> loop xs
+              _ -> mustThrowPyError
+      loop $ Set.toList set
+
+instance (FromPy a, Ord a) => FromPy (Set.Set a) where
+  basicFromPy p_set = basicGetIter p_set >>= \case
+    NULL -> do Py [C.exp| void { PyErr_Clear() } |]
+               throwM BadPyType
+    p_iter -> foldPyIterable p_iter
+      (\s p -> do a <- basicFromPy p
+                  pure $! Set.insert a s)
+      Set.empty
+
+-- | Fold over iterable. Function takes ownership over iterator.
+foldPyIterable
+  :: Ptr PyObject                -- ^ Python iterator (not checked)
+  -> (a -> Ptr PyObject -> Py a) -- ^ Step function. It takes borrowed pointer.
+  -> a                           -- ^ Initial value
+  -> Py a
+foldPyIterable p_iter step a0
+  = loop a0 `finally` decref p_iter
+  where
+    loop a = basicIterNext p_iter >>= \case
+      NULL -> a <$ checkThrowPyError
+      p    -> loop =<< (step a p `finally` decref p)
 
 
 ----------------------------------------------------------------
@@ -436,7 +466,7 @@ instance (ToPy b) => ToPy (IO b) where
   basicToPy f = Py $ do
     --
     f_ptr <- wrapCFunction $ \_ _ -> pyCallback $ do
-      lift $ basicToPy =<< dropGIL f
+      progPy $ basicToPy =<< dropGIL f
     --
     [CU.exp| PyObject* { inline_py_callback_METH_NOARGS($(PyCFunction f_ptr)) } |]
 
@@ -447,7 +477,7 @@ instance (FromPy a, Show a, ToPy b) => ToPy (a -> IO b) where
     --
     f_ptr <- wrapCFunction $ \_ p_a -> pyCallback $ do
       a <- loadArg p_a 0 1
-      lift $ basicToPy =<< dropGIL (f a)
+      progPy $ basicToPy =<< dropGIL (f a)
     --
     [CU.exp| PyObject* { inline_py_callback_METH_O($(PyCFunction f_ptr)) } |]
 
@@ -459,7 +489,7 @@ instance (FromPy a1, FromPy a2, ToPy b) => ToPy (a1 -> a2 -> IO b) where
       when (n /= 2) $ abortM $ raiseBadNArgs 2 n
       a1 <- loadArgFastcall p_arr 0 n
       a2 <- loadArgFastcall p_arr 1 n
-      lift $ basicToPy =<< dropGIL (f a1 a2)
+      progPy $ basicToPy =<< dropGIL (f a1 a2)
     --
     [CU.exp| PyObject* { inline_py_callback_METH_FASTCALL($(PyCFunctionFast f_ptr)) } |]
 
@@ -470,7 +500,7 @@ instance (FromPy a1, FromPy a2, ToPy b) => ToPy (a1 -> a2 -> IO b) where
 
 -- | Execute haskell callback function
 pyCallback :: Program (Ptr PyObject) (Ptr PyObject) -> IO (Ptr PyObject)
-pyCallback io = callbackEnsurePyLock $ unPy $ ensureGIL $ evalContT io `catch` convertHaskell2Py
+pyCallback io = callbackEnsurePyLock $ unPy $ ensureGIL $ runProgram io `catch` convertHaskell2Py
 
 -- | Load argument from python object for haskell evaluation
 loadArg
@@ -479,7 +509,7 @@ loadArg
   -> Int            -- ^ Argument number (0-based)
   -> Int64          -- ^ Total number of arguments
   -> Program (Ptr PyObject) a
-loadArg p (fromIntegral -> i) (fromIntegral -> tot) = ContT $ \success -> do
+loadArg p (fromIntegral -> i) (fromIntegral -> tot) = Program $ ContT $ \success -> do
   try (basicFromPy p) >>= \case
     Right a          -> success a
     Left  BadPyType  -> oops
@@ -501,7 +531,7 @@ loadArgFastcall
   -> Int64              -- ^ Total number of arguments
   -> Program (Ptr PyObject) a
 loadArgFastcall p_arr i tot = do
-  p <- liftIO $ peekElemOff p_arr i
+  p <- progIO $ peekElemOff p_arr i
   loadArg p i tot
 
 raiseBadNArgs :: CInt -> Int64 -> Py (Ptr PyObject)

src/Python/Internal/CAPI.hs

@@ -0,0 +1,58 @@
+{-# LANGUAGE QuasiQuotes     #-}
+{-# LANGUAGE TemplateHaskell #-}
+-- |
+-- Thin wrappers over C API
+module Python.Internal.CAPI
+  ( decref
+  , incref
+    -- * Simple wrappers
+  , basicNewDict
+  , basicNewSet
+  , basicGetIter
+  , basicIterNext
+  , basicCallKwdOnly
+  ) where
+
+import Foreign.Ptr
+import Language.C.Inline          qualified as C
+import Language.C.Inline.Unsafe   qualified as CU
+
+import Python.Internal.Types
+
+
+----------------------------------------------------------------
+C.context (C.baseCtx <> pyCtx)
+C.include "<inline-python.h>"
+----------------------------------------------------------------
+
+
+decref :: Ptr PyObject -> Py ()
+decref p = Py [CU.exp| void { Py_DECREF($(PyObject* p)) } |]
+
+incref :: Ptr PyObject -> Py ()
+incref p = Py [CU.exp| void { Py_INCREF($(PyObject* p)) } |]
+
+basicNewDict :: Py (Ptr PyObject)
+basicNewDict = Py [CU.exp| PyObject* { PyDict_New() } |]
+
+basicNewSet :: Py (Ptr PyObject)
+basicNewSet = Py [CU.exp| PyObject* { PySet_New(NULL) } |]
+
+basicGetIter :: Ptr PyObject -> Py (Ptr PyObject)
+basicGetIter p = Py [CU.exp| PyObject* { PyObject_GetIter( $(PyObject *p)) } |]
+
+basicIterNext :: Ptr PyObject -> Py (Ptr PyObject)
+basicIterNext p = Py [C.exp| PyObject* { PyIter_Next($(PyObject* p)) } |]
+
+
+-- | Call python function using only keyword arguments
+basicCallKwdOnly
+  :: Ptr PyObject -- ^ Function object
+  -> Ptr PyObject -- ^ Keywords. Must be dictionary
+  -> Py (Ptr PyObject)
+basicCallKwdOnly fun kwd = Py [CU.block| PyObject* {
+  PyObject* args = PyTuple_Pack(0);
+  PyObject* res  = PyObject_Call($(PyObject *fun), args, $(PyObject *kwd));
+  Py_DECREF(args);
+  return res;
+  } |]