Add instances for ShortByteString

Author: Shimuuar

src/Python/Inline/Literal.hs

@@ -14,6 +14,7 @@ module Python.Inline.Literal
   , fromPy'
   ) where
 
+import Control.Exception           (evaluate)
 import Control.Monad
 import Control.Monad.Catch
 import Control.Monad.Trans.Cont
@@ -23,6 +24,7 @@ import Data.Int
 import Data.Word
 import Data.ByteString             qualified as BS
 import Data.ByteString.Unsafe      qualified as BS
+import Data.ByteString.Short       qualified as SBS
 import Data.ByteString.Lazy        qualified as BL
 import Data.Set                    qualified as Set
 import Data.Map.Strict             qualified as Map
@@ -577,6 +579,36 @@ instance ToPy BL.ByteString where
 instance FromPy BL.ByteString where
   basicFromPy = fmap BL.fromStrict . basicFromPy
 
+
+-- | @since NEXT_VERSION@. Accepts @bytes@ and @bytearray@
+instance FromPy SBS.ShortByteString where
+  basicFromPy py = pyIO $ do
+    [CU.exp| int { PyBytes_Check($(PyObject* py)) } |] >>= \case
+      TRUE -> do
+        sz  <- [CU.exp| int64_t { PyBytes_GET_SIZE( $(PyObject* py)) } |]
+        buf <- [CU.exp| char*   { PyBytes_AS_STRING($(PyObject* py)) } |]
+        fini buf (fromIntegral sz)
+      _ -> [CU.exp| int { PyByteArray_Check($(PyObject* py)) } |] >>= \case
+        TRUE -> do
+          sz  <- [CU.exp| int64_t { PyByteArray_GET_SIZE( $(PyObject* py)) } |]
+          buf <- [CU.exp| char*   { PyByteArray_AS_STRING($(PyObject* py)) } |]
+          fini buf (fromIntegral sz)
+        _ -> throwM BadPyType
+    where
+      fini buf sz = do
+        bs <- BS.unsafePackCStringLen (buf, sz)
+        evaluate $ SBS.toShort bs
+
+-- | @since NEXT_VERSION@. Converted to @bytes@
+instance ToPy SBS.ShortByteString where
+  basicToPy bs = pyIO $ SBS.useAsCStringLen bs $ \(ptr,len) -> do
+    let c_len = fromIntegral len :: CLLong
+    py <- [CU.exp| PyObject* { PyBytes_FromStringAndSize($(char* ptr), $(long long c_len)) }|]
+    case py of
+      NULL -> unsafeRunPy mustThrowPyError
+      _    -> return py
+
+
 -- | @since NEXT_VERSION@.
 instance ToPy T.Text where
   -- NOTE: Is there ore efficient way to access

test/TST/Roundtrip.hs

@@ -22,6 +22,7 @@ import Python.Inline.QQ
 
 import Data.ByteString             qualified as BS
 import Data.ByteString.Lazy        qualified as BL
+import Data.ByteString.Short       qualified as SBS
 import Data.Vector                 qualified as V
 #if MIN_VERSION_vector(0,13,2)
 import Data.Vector.Strict          qualified as VV
@@ -82,6 +83,7 @@ tests = testGroup "Roundtrip"
 #endif
     , testRoundtrip @BS.ByteString
     , testRoundtrip @BL.ByteString
+    , testRoundtrip @SBS.ShortByteString
     , testRoundtrip @T.Text
     , testRoundtrip @TL.Text
     ]