Context: adding ability to add custom load callback

steweg · steweg · commit 7f42a373db48 · 2024-02-28T13:12:33.000+01:00
This patch adds ability to add custom module load callback, which
allows user to load modules from remote source etc.
diff --git a/cffi/cdefs.h b/cffi/cdefs.h
@@ -935,6 +935,11 @@ typedef enum {
 LY_ERR lys_parse(struct ly_ctx *, struct ly_in *, LYS_INFORMAT, const char **, struct lys_module **);
 LY_ERR ly_ctx_new_ylpath(const char *, const char *, LYD_FORMAT, int, struct ly_ctx **);
 LY_ERR ly_ctx_get_yanglib_data(const struct ly_ctx *, struct lyd_node **, const char *, ...);
+typedef void (*ly_module_imp_data_free_clb)(void *, void *);
+typedef LY_ERR (*ly_module_imp_clb)(const char *, const char *, const char *, const char *, void *, LYS_INFORMAT *, const char **, ly_module_imp_data_free_clb *);
+void ly_ctx_set_module_imp_clb(struct ly_ctx *, ly_module_imp_clb, void *);
+extern "Python" void lypy_module_imp_data_free_clb(void *, void *);
+extern "Python" LY_ERR lypy_module_imp_clb(const char *, const char *, const char *, const char *, void *, LYS_INFORMAT *, const char **, ly_module_imp_data_free_clb *);
 
 struct lyd_meta {
     struct lyd_node *parent;
diff --git a/libyang/context.py b/libyang/context.py
@@ -4,7 +4,7 @@
 # SPDX-License-Identifier: MIT
 
 import os
-from typing import IO, Any, Iterator, Optional, Union
+from typing import IO, Any, Callable, Iterator, Optional, Tuple, Union
 
 from _libyang import ffi, lib
 from .data import (
@@ -19,9 +19,48 @@
 from .util import DataType, IOType, LibyangError, c2str, data_load, str2c
 
 
+# -------------------------------------------------------------------------------------
+@ffi.def_extern(name="lypy_module_imp_data_free_clb")
+def libyang_c_module_imp_data_free_clb(cdata, user_data):
+    instance = ffi.from_handle(user_data)
+    instance.free_module_data(cdata)
+
+
+# -------------------------------------------------------------------------------------
+@ffi.def_extern(name="lypy_module_imp_clb")
+def libyang_c_module_imp_clb(
+    mod_name,
+    mod_rev,
+    submod_name,
+    submod_rev,
+    user_data,
+    fmt,
+    module_data,
+    free_module_data,
+):
+    fmt[0] = lib.LYS_IN_UNKNOWN
+    module_data[0] = ffi.NULL
+    free_module_data[0] = lib.lypy_module_imp_data_free_clb
+    instance = ffi.from_handle(user_data)
+    in_fmt, content = instance.get_module_data(
+        c2str(mod_name), c2str(mod_rev), c2str(submod_name), c2str(submod_rev)
+    )
+    if content is None:
+        return lib.LY_ENOT
+    fmt[0] = schema_in_format(in_fmt)
+    module_data[0] = content
+    return lib.LY_SUCCESS
+
+
 # -------------------------------------------------------------------------------------
 class Context:
-    __slots__ = ("cdata", "__dict__")
+    __slots__ = (
+        "cdata",
+        "_module_data_clb",
+        "_cffi_handle",
+        "_cdata_modules",
+        "__dict__",
+    )
 
     def __init__(
         self,
@@ -470,3 +509,41 @@ def __iter__(self) -> Iterator[Module]:
         while mod:
             yield Module(self, mod)
             mod = lib.ly_ctx_get_module_iter(self.cdata, idx)
+
+    def free_module_data(self, cdata) -> None:
+        self._cdata_modules.remove(cdata)
+
+    def get_module_data(
+        self,
+        mod_name: Optional[str],
+        mod_rev: Optional[str],
+        submod_name: Optional[str],
+        submod_rev: Optional[str],
+    ) -> Tuple[str, Optional[str]]:
+        if self._module_data_clb is None:
+            return None
+        fmt_str, module_data = self._module_data_clb(
+            mod_name, mod_rev, submod_name, submod_rev
+        )
+        if module_data is None:
+            return fmt_str, None
+        module_data_c = str2c(module_data)
+        self._cdata_modules.append(module_data_c)
+        return fmt_str, module_data_c
+
+    def set_module_data_clb(
+        self,
+        clb: Optional[
+            Callable[
+                [Optional[str], Optional[str], Optional[str], Optional[str]],
+                Tuple[str, Optional[str]],
+            ]
+        ] = None,
+    ) -> None:
+        self._module_data_clb = clb
+        if clb is None:
+            lib.ly_ctx_set_module_imp_clb(self.cdata, ffi.NULL, ffi.NULL)
+        else:
+            lib.ly_ctx_set_module_imp_clb(
+                self.cdata, lib.lypy_module_imp_clb, self._cffi_handle
+            )
diff --git a/tests/test_context.py b/tests/test_context.py
@@ -116,3 +116,18 @@ def test_ctx_disable_searchdirs(self):
         with Context(YANG_DIR, disable_searchdirs=True) as ctx:
             with self.assertRaises(LibyangError):
                 ctx.load_module("yolo-nodetypes")
+
+    def test_ctx_using_clb(self):
+        def get_module_clb(mod_name, *_):
+            YOLO_NODETYPES_MOD_PATH = os.path.join(YANG_DIR, "yolo/yolo-nodetypes.yang")
+            self.assertEqual(mod_name, "yolo-nodetypes")
+            with open(YOLO_NODETYPES_MOD_PATH, encoding="utf-8") as f:
+                mod_str = f.read()
+            return "yang", mod_str
+
+        with Context(YANG_DIR, disable_searchdirs=True) as ctx:
+            with self.assertRaises(LibyangError):
+                ctx.load_module("yolo-nodetypes")
+            ctx.set_module_data_clb(get_module_clb)
+            mod = ctx.load_module("yolo-nodetypes")
+            self.assertIsInstance(mod, Module)
