Bug 2374 - YANG Binding: Added support for AugmentationHolder interface

[mdsal.git] / yang / yang-binding / src / main / java / org / opendaylight / yangtools / yang / binding / util / AugmentationFieldGetter.java

diff --git a/yang/yang-binding/src/main/java/org/opendaylight/yangtools/yang/binding/util/AugmentationFieldGetter.java b/yang/yang-binding/src/main/java/org/opendaylight/yangtools/yang/binding/util/AugmentationFieldGetter.java
index 8178d8f1e1923580d490028cd6222159b1e69ace..2e8a33a272b7f542ba95a731679833f73914e777 100644 (file)
--- a/yang/yang-binding/src/main/java/org/opendaylight/yangtools/yang/binding/util/AugmentationFieldGetter.java
+++ b/yang/yang-binding/src/main/java/org/opendaylight/yangtools/yang/binding/util/AugmentationFieldGetter.java
@@ -1,12 +1,14 @@
 /*
- * Copyright (c) 2014 Cisco Systems, Inc. and others.  All rights reserved.
+ * Copyright (c) 2014 Cisco Systems, Inc. and others. All rights reserved.
  *
- * This program and the accompanying materials are made available under the
- * terms of the Eclipse Public License v1.0 which accompanies this distribution,
- * and is available at http://www.eclipse.org/legal/epl-v10.html
+ * This program and the accompanying materials are made available under the terms of the Eclipse
+ * Public License v1.0 which accompanies this distribution, and is available at
+ * http://www.eclipse.org/legal/epl-v10.html
  */
 package org.opendaylight.yangtools.yang.binding.util;
 
+import org.opendaylight.yangtools.yang.binding.AugmentationHolder;
+
 import com.google.common.base.Preconditions;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
@@ -30,19 +32,31 @@ abstract class AugmentationFieldGetter {
         }
     };
 
-   /**
-    *
-    * Retrieves augmentations from supplied object
-    *
-    * @param input Input Data object, from which augmentations should be extracted
-    * @return Map of Augmentation class to augmentation
-    */
-   protected abstract Map<Class<? extends Augmentation<?>>, Augmentation<?>> getAugmentations(final Object input);
+    private static final AugmentationFieldGetter AUGMENTATION_HOLDER_GETTER = new AugmentationFieldGetter() {
 
-    private static final LoadingCache<Class<?>, AugmentationFieldGetter> AUGMENTATION_GETTERS =
-            CacheBuilder.newBuilder().weakKeys().softValues().build(new AugmentationGetterLoader());
+        @Override
+        @SuppressWarnings({"unchecked", "rawtypes"})
+        protected Map<Class<? extends Augmentation<?>>, Augmentation<?>> getAugmentations(final Object input) {
+            return (Map) ((AugmentationHolder<?>) input).augmentations();
+        }
+    };
+
+    /**
+     *
+     * Retrieves augmentations from supplied object
+     *
+     * @param input Input Data object, from which augmentations should be extracted
+     * @return Map of Augmentation class to augmentation
+     */
+    protected abstract Map<Class<? extends Augmentation<?>>, Augmentation<?>> getAugmentations(final Object input);
+
+    private static final LoadingCache<Class<?>, AugmentationFieldGetter> AUGMENTATION_GETTERS = CacheBuilder
+            .newBuilder().weakKeys().softValues().build(new AugmentationGetterLoader());
 
     public static AugmentationFieldGetter getGetter(final Class<? extends Object> clz) {
+        if(AugmentationHolder.class.isAssignableFrom(clz)) {
+            return AUGMENTATION_HOLDER_GETTER;
+        }
         return AUGMENTATION_GETTERS.getUnchecked(clz);
     }
 
@@ -81,5 +95,4 @@ abstract class AugmentationFieldGetter {
         }
     }
 
-
 }