From adab39c83947c9c08d3a39d212f687f002b8c13e Mon Sep 17 00:00:00 2001
From: Stefan Gula <steweg@gmail.com>
Date: Mon, 5 Aug 2024 12:25:42 +0200
Subject: [PATCH] schema: adding parent_node API to extensions

This patch adds parent_node API, which allows caller to get parent PNode
from PExtension or SNode from Extension instance.

Signed-off-by: Stefan Gula <steweg@gmail.com>
Signed-off-by: Samuel Gauthier <samuel.gauthier@6wind.com>
---
 cffi/cdefs.h         |  2 ++
 libyang/__init__.py  |  1 +
 libyang/schema.py    | 16 ++++++++++++++++
 tests/test_schema.py |  7 +++++++
 4 files changed, 26 insertions(+)

diff --git a/cffi/cdefs.h b/cffi/cdefs.h
index 40d3669..1c1d8f3 100644
--- a/cffi/cdefs.h
+++ b/cffi/cdefs.h
@@ -174,6 +174,8 @@ enum ly_stmt {
     LY_STMT_ARG_VALUE
 };
 
+#define LY_STMT_NODE_MASK ...
+
 #define LY_LOLOG ...
 #define LY_LOSTORE ...
 #define LY_LOSTORE_LAST ...
diff --git a/libyang/__init__.py b/libyang/__init__.py
index 26235de..cab9971 100644
--- a/libyang/__init__.py
+++ b/libyang/__init__.py
@@ -67,6 +67,7 @@
 from .log import configure_logging
 from .schema import (
     Extension,
+    ExtensionParsed,
     Feature,
     IfAndFeatures,
     IfFeature,
diff --git a/libyang/schema.py b/libyang/schema.py
index f35a5e9..a47974e 100644
--- a/libyang/schema.py
+++ b/libyang/schema.py
@@ -411,6 +411,14 @@ def name(self) -> str:
     def module(self) -> Module:
         return self._module_from_parsed()
 
+    def parent_node(self) -> Optional["PNode"]:
+        if not bool(self.cdata.parent_stmt & lib.LY_STMT_NODE_MASK):
+            return None
+        try:
+            return PNode.new(self.context, self.cdata.parent, self.module())
+        except LibyangError:
+            return None
+
 
 # -------------------------------------------------------------------------------------
 class ExtensionCompiled(Extension):
@@ -428,6 +436,14 @@ def module(self) -> Module:
             raise self.context.error("cannot get module")
         return Module(self.context, self.cdata_def.module)
 
+    def parent_node(self) -> Optional["SNode"]:
+        if not bool(self.cdata.parent_stmt & lib.LY_STMT_NODE_MASK):
+            return None
+        try:
+            return SNode.new(self.context, self.cdata.parent)
+        except LibyangError:
+            return None
+
 
 # -------------------------------------------------------------------------------------
 class _EnumBit:
diff --git a/tests/test_schema.py b/tests/test_schema.py
index e2a833c..1ae9fdf 100644
--- a/tests/test_schema.py
+++ b/tests/test_schema.py
@@ -7,6 +7,7 @@
 from libyang import (
     Context,
     Extension,
+    ExtensionParsed,
     IfFeature,
     IfOrFeatures,
     IOType,
@@ -496,6 +497,11 @@ def test_rpc_extensions(self):
         self.assertEqual(len(ext), 1)
         ext = self.rpc.get_extension("require-admin", prefix="omg-extensions")
         self.assertIsInstance(ext, Extension)
+        self.assertIsInstance(ext.parent_node(), SRpc)
+        parsed = self.rpc.parsed()
+        ext = parsed.get_extension("require-admin", prefix="omg-extensions")
+        self.assertIsInstance(ext, ExtensionParsed)
+        self.assertIsInstance(ext.parent_node(), PAction)
 
     def test_rpc_params(self):
         leaf = next(self.rpc.children())
@@ -609,6 +615,7 @@ def test_leaf_type_extensions(self):
             "type-desc", prefix="omg-extensions", arg_value="<protocol>"
         )
         self.assertIsInstance(ext, Extension)
+        self.assertIsNone(ext.parent_node())
 
     def test_leaf_type_enum(self):
         leaf = next(