check for transform_reply type (#270)

ncclient · Aug 25, 2018 · df3311c · df3311c
1 parent a475d87
commit df3311c
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/ncclient/xml_.py b/ncclient/xml_.py
@@ -19,6 +19,7 @@
 import io
 import sys
 import six
+import types
 from six import StringIO
 from io import BytesIO
 from lxml import etree
@@ -148,12 +149,15 @@ def validated_element(x, tags=None, attrs=None):
     're':'http://exslt.org/regular-expressions'
 }
 
+
 class NCElement(object):
     def __init__(self, result, transform_reply):
         self.__result = result
         self.__transform_reply = transform_reply
-        self.__doc = self.remove_namespaces(self.__result)
-
+        if isinstance(transform_reply, types.FunctionType):
+            self.__doc = self.__transform_reply(result._root)
+        else:
+            self.__doc = self.remove_namespaces(self.__result)
 
     def xpath(self, expression):
         """