From f35afca564d85ae01f5fd06858f5e00d9a8e993e Mon Sep 17 00:00:00 2001
From: Robert Varga <robert.varga@pantheon.tech>
Date: Tue, 20 Jun 2023 03:08:25 +0200
Subject: [PATCH] Add
 BindingDataObjectCodecTreeParent.getStream{Augmentation,DataObject}

When we know we are streaming to a particular type, we can improve the
user experience by returning a friendly interface.

JIRA: MDSAL-820
Change-Id: Iadb85de866a0446934f636bb870e1df5f4c81a4b
Signed-off-by: Robert Varga <robert.varga@pantheon.tech>
---
 .../api/BindingDataObjectCodecTreeParent.java | 21 ++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/binding/mdsal-binding-dom-codec-api/src/main/java/org/opendaylight/mdsal/binding/dom/codec/api/BindingDataObjectCodecTreeParent.java b/binding/mdsal-binding-dom-codec-api/src/main/java/org/opendaylight/mdsal/binding/dom/codec/api/BindingDataObjectCodecTreeParent.java
index e7cb9ddea0..bfaf2aba4a 100644
--- a/binding/mdsal-binding-dom-codec-api/src/main/java/org/opendaylight/mdsal/binding/dom/codec/api/BindingDataObjectCodecTreeParent.java
+++ b/binding/mdsal-binding-dom-codec-api/src/main/java/org/opendaylight/mdsal/binding/dom/codec/api/BindingDataObjectCodecTreeParent.java
@@ -9,6 +9,7 @@ package org.opendaylight.mdsal.binding.dom.codec.api;
 
 import com.google.common.annotations.Beta;
 import org.eclipse.jdt.annotation.NonNull;
+import org.opendaylight.yangtools.yang.binding.Augmentation;
 import org.opendaylight.yangtools.yang.binding.ChoiceIn;
 import org.opendaylight.yangtools.yang.binding.DataObject;
 
@@ -37,7 +38,7 @@ import org.opendaylight.yangtools.yang.binding.DataObject;
 public interface BindingDataObjectCodecTreeParent<T> {
     /**
      * Returns child context as if it was walked by {@link BindingStreamEventWriter}. This means that to enter case,
-     * one must issue {@code streamChild(ChoiceClass).streamChild(CaseClass)}.
+     * one must issue {@code getStreamChild(ChoiceClass).getStreamChild(CaseClass)}.
      *
      * @param <E> Stream child DataObject type
      * @param childClass Child class by Binding Stream navigation
@@ -46,4 +47,22 @@ public interface BindingDataObjectCodecTreeParent<T> {
      * @throws IllegalArgumentException If supplied child class is not valid in specified context.
      */
     <E extends DataObject> @NonNull CommonDataObjectCodecTreeNode<E> getStreamChild(@NonNull Class<E> childClass);
+
+    default <A extends Augmentation<?>> @NonNull BindingAugmentationCodecTreeNode<A> getStreamAugmentation(
+            final @NonNull Class<A> childClass) {
+        final var result = getStreamChild(childClass);
+        if (result instanceof BindingAugmentationCodecTreeNode) {
+            return (BindingAugmentationCodecTreeNode<A>) result;
+        }
+        throw new IllegalArgumentException("Child " + childClass.getName() + " results in non-Augmentation " + result);
+    }
+
+    default <E extends DataObject> @NonNull BindingDataObjectCodecTreeNode<E> getStreamDataObject(
+            final @NonNull Class<E> childClass) {
+        final var result = getStreamChild(childClass);
+        if (result instanceof BindingDataObjectCodecTreeNode) {
+            return (BindingDataObjectCodecTreeNode<E>) result;
+        }
+        throw new IllegalArgumentException("Child " + childClass.getName() + " results in non-DataObject " + result);
+    }
 }
-- 
2.36.6