From 0ff07a79e472598754db534abc342cfc738dc9d3 Mon Sep 17 00:00:00 2001
From: Gabriele Tornetta <gabriele.tornetta@datadoghq.com>
Date: Fri, 14 Oct 2022 11:38:32 +0100
Subject: [PATCH] use c_long as default for argument and return type

---
 test/cunit/__init__.py   | 33 ++++++++++++++++++++-------------
 test/cunit/test_cache.py |  4 +++-
 2 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/test/cunit/__init__.py b/test/cunit/__init__.py
index 10398aae..9d27c403 100644
--- a/test/cunit/__init__.py
+++ b/test/cunit/__init__.py
@@ -1,10 +1,10 @@
 import ctypes
 import re
-from ctypes import CDLL, POINTER, Structure, c_char_p, c_void_p, cast
+from ctypes import CDLL, POINTER, Structure, c_char_p, c_long, c_void_p, cast
 from pathlib import Path
 from subprocess import PIPE, STDOUT, run
 from types import ModuleType
-from typing import Any, Callable, Optional, Type
+from typing import Any, Callable, Optional, Type, Union
 
 from pycparser import c_ast, c_parser
 from pycparser.plyparser import ParseError
@@ -88,13 +88,11 @@ def compile(
 
 C = CDLL("libc.so.6")
 
-# https://github.com/actions/setup-python/issues/442#issuecomment-1193140956
-C.malloc.restype = ctypes.c_void_p
-C.free.argtypes = [ctypes.c_void_p]
-
 
 class CFunctionDef:
-    def __init__(self, name: str, args: list[str], rtype: Any) -> None:
+    def __init__(
+        self, name: str, args: list[tuple[str, Union[c_long, c_void_p]]], rtype: Any
+    ) -> None:
         self.name = name
         self.args = args
         self.rtype = rtype
@@ -117,16 +115,20 @@ def __del__(self) -> None:
             self.destroy()
 
     def __repr__(self) -> str:
-        return f"<{self.name} CObject at {self.__cself__}>"
+        return f"<{self.__class__.__name__} CObject at {self.__cself__}>"
 
 
 class CFunction:
     def __init__(self, cfuncdef: CFunctionDef, cfunc: Callable[..., Any]) -> None:
         self.__name__ = cfuncdef.name
-        self.__args__ = cfuncdef.args
+        self.__args__ = [_[0] for _ in cfuncdef.args]
         self.__cfunc__ = cfunc
-        if cfuncdef.rtype is not None:
-            self.__cfunc__.restype = cfuncdef.rtype
+
+        # Prevent argument values from being truncated/mangled
+        self.__cfunc__.argtypes = [_[1] for _ in cfuncdef.args]
+        self.__cfunc__.restype = (
+            cfuncdef.rtype if cfuncdef.rtype is not None else c_long
+        )
 
         self._posonly = all(_ is None for _ in self.__args__)
 
@@ -211,7 +213,6 @@ def __init__(self) -> None:
         self.functions = []
 
     def _get_type(self, node: c_ast.Node) -> None:
-        print(node)
         return self.types[" ".join(node.type.type.names)]
 
     def visit_Typedef(self, node: c_ast.Node) -> None:
@@ -237,8 +238,14 @@ def visit_Decl(self, node: c_ast.Node) -> None:
                     rtype = c_char_p
                 else:
                     rtype = c_void_p
+
             args = (
-                [_.name if hasattr(_, "name") else None for _ in node.type.args.params]
+                [
+                    (_.name, c_void_p if isinstance(_.type, c_ast.PtrDecl) else c_long)
+                    if hasattr(_, "name")
+                    else None
+                    for _ in node.type.args.params
+                ]
                 if node.type.args is not None
                 else []
             )
diff --git a/test/cunit/test_cache.py b/test/cunit/test_cache.py
index ec809bd1..08e15f79 100644
--- a/test/cunit/test_cache.py
+++ b/test/cunit/test_cache.py
@@ -1,9 +1,11 @@
+from ctypes import c_void_p
 from test.cunit import C
 from test.cunit.cache import Chain, HashTable, LruCache, Queue, QueueItem
 
 import pytest
 
 NULL = 0
+C.free.argtypes = [c_void_p]
 
 
 def test_queue_item():
@@ -31,7 +33,7 @@ def test_queue(qsize):
 
     assert qsize == 0 or not q.is_empty()
     assert q.is_full()
-    assert q.enqueue(42, 42) is NULL
+    assert q.enqueue(42, 42) is None
 
     assert values == [q.dequeue() for _ in range(qsize)]