Add optional libraries

xobjects/context.py

@@ -354,25 +354,75 @@ def build_kernels(
     ) -> Dict[Tuple[str, tuple], KernelType]:
         pass
 
-    def get_installed_c_source_paths(self) -> List[str]:
-        """Returns a list of include paths registered in dependent packages.
+    def get_installed_c_source_and_library_paths(
+        self,
+    ) -> tuple[set[Path], set[str], set[Path]]:
+        """Returns a list of C paths registered in dependent packages.
 
-        In a package that depends on xobjects, you can register C source paths
-        using the entry point `xobjects.c_sources`. A path to the directory
-        containing the specified module will be added to the include path when
-        building kernels. For example, the following will allow to write
-        ``#include <xtrack/path/to/some/header.h>`` in kernel sources:
+        In a package that depends on xobjects, you can register C source and
+        library paths using the entry point `xobjects.build_info`. These paths
+        will be added to the C include path and the library path when building
+        kernels. For example, the following will allow to write
+        ``#include <xcoll/path/to/some/header.h>`` in kernel sources, and
+        allow to use functions from the library ``xcoll/lib/libFlukaIO.a``:
 
          .. code-block:: toml
             [project.entry-points.xobjects]
-            include = "xtrack"
+            build_info = "xcoll._xobjects:get_build_info"
+
+        and in the file ``xcoll/_xobjects.py``:
+
+         .. code-block:: python
+            from ..general import _pkg_root
+            def get_build_info():
+                return {
+                    "include_dirs": [_pkg_root.parent],
+                    "libraries": ["FlukaIO"],
+                    "library_dirs": [_pkg_root / "lib"],
+                }
         """
-        sources = []
+        sources = set()
+        libs = set()
+        lib_paths = set()
+
+        # Old entry point for backward compatibility
         for ep in entry_points(group="xobjects", name="include"):
             module = ep.load()
             path = Path(module.__file__).parents[1]
-            sources.append(str(path))
-        return sources
+            sources.add(path)
+
+        # New entry point
+        for ep in entry_points(group="xobjects", name="build_info"):
+            get_build_info = ep.load()
+            info = get_build_info()
+            include_dirs = info.get("include_dirs", [])
+            if not hasattr(include_dirs, "__iter__") or isinstance(
+                include_dirs, str
+            ):
+                include_dirs = [include_dirs]
+            include_dirs = [
+                Path(dd).expanduser().resolve() for dd in include_dirs
+            ]
+            sources.update(include_dirs)
+
+            library_dirs = info.get("library_dirs", [])
+            if not hasattr(library_dirs, "__iter__") or isinstance(
+                library_dirs, str
+            ):
+                library_dirs = [library_dirs]
+            library_dirs = [
+                Path(dd).expanduser().resolve() for dd in library_dirs
+            ]
+            lib_paths.update(library_dirs)
+
+            libraries = info.get("libraries", [])
+            if not hasattr(libraries, "__iter__") or isinstance(
+                libraries, str
+            ):
+                libraries = [libraries]
+            libs.update(libraries)
+
+        return sources, libs, lib_paths
 
     @abstractmethod
     def nparray_to_context_array(self, arr, copy=False):

xobjects/context_cpu.py

@@ -499,10 +499,11 @@ def compile_kernel(
             xtr_compile_args.append("-DXO_CONTEXT_CPU_SERIAL")
             xtr_link_args.append("-DXO_CONTEXT_CPU_SERIAL")
 
-        extra_include_paths = self.get_installed_c_source_paths()
-        include_flags = [f"-I{path}" for path in extra_include_paths]
-        xtr_compile_args.extend(include_flags)
-        xtr_link_args.extend(include_flags)
+        (
+            extra_include_paths,
+            extra_libraries,
+            extra_library_paths,
+        ) = self.get_installed_c_source_and_library_paths()
 
         if os.name == "nt":  # windows
             # TODO: to be handled properly
@@ -516,6 +517,9 @@ def compile_kernel(
         ffi_interface.set_source(
             module_name,
             specialized_source,
+            include_dirs=[path.as_posix() for path in extra_include_paths],
+            libraries=list(extra_libraries),
+            library_dirs=[path.as_posix() for path in extra_library_paths],
             extra_compile_args=xtr_compile_args,
             extra_link_args=xtr_link_args,
         )

xobjects/context_cupy.py

@@ -356,7 +356,8 @@ def __invert__(self):
             return cupy.ndarray.__invert__(self._as_cupy())
 
 
-cudaheader: List[SourceType] = ["""\
+cudaheader: List[SourceType] = [
+    """\
 typedef signed int         int32_t;  //only_for_context cuda
 typedef signed short       int16_t;  //only_for_context cuda
 typedef signed char        int8_t;   //only_for_context cuda
@@ -373,7 +374,8 @@ def __invert__(self):
   #define NULL nullptr
 #endif
 
-"""]
+"""
+]
 
 
 def nplike_to_cupy(arr):
@@ -477,8 +479,15 @@ def build_kernels(
             with open(save_source_as, "w") as fid:
                 fid.write(specialized_source)
 
-        extra_include_paths = self.get_installed_c_source_paths()
-        include_flags = [f"-I{path}" for path in extra_include_paths]
+        (
+            # TODO: how to deal with CUDA libraries?
+            extra_include_paths,
+            _,
+            _,
+        ) = self.get_installed_c_source_and_library_paths()
+        include_flags = [
+            f"-I{path.as_posix()}" for path in extra_include_paths
+        ]
         extra_compile_args = (
             *extra_compile_args,
             *include_flags,

xobjects/context_pyopencl.py

@@ -44,7 +44,8 @@
 
 from ._patch_pyopencl_array import _patch_pyopencl_array
 
-openclheader: List[SourceType] = ["""\
+openclheader: List[SourceType] = [
+    """\
 #ifndef XOBJ_STDINT
 typedef long           int64_t;
 typedef int            int32_t;
@@ -58,7 +59,8 @@
 #ifndef NULL
 #define NULL 0L
 #endif
-"""]
+"""
+]
 
 if _enabled:
     # order of base classes matters as it defines which __setitem__ is used
@@ -228,8 +230,15 @@ def build_kernels(
             with open(save_source_as, "w") as fid:
                 fid.write(specialized_source)
 
-        extra_include_paths = self.get_installed_c_source_paths()
-        include_flags = [f"-I{path}" for path in extra_include_paths]
+        (
+            # TODO: how to deal with OpenCL libraries?
+            extra_include_paths,
+            _,
+            _,
+        ) = self.get_installed_c_source_and_library_paths()
+        include_flags = [
+            f"-I{path.as_posix()}" for path in extra_include_paths
+        ]
 
         extra_compile_args = (
             *extra_compile_args,