From 10def22633f0f736dd955f071c353af8ba3e7199 Mon Sep 17 00:00:00 2001
From: jmarinaro <jmarinaro@users.noreply.github.com>
Date: Fri, 14 Jun 2019 15:00:45 -0600
Subject: [PATCH] improved retcode handling Fixes #195

---
 pepper/retcode.py | 77 +++++++++++++++++++++++------------------------
 1 file changed, 38 insertions(+), 39 deletions(-)

diff --git a/pepper/retcode.py b/pepper/retcode.py
index 11b254c..12f8f02 100644
--- a/pepper/retcode.py
+++ b/pepper/retcode.py
@@ -29,6 +29,29 @@ def validate(self, options, result):
             return self.validate_fail_all_none(result)
         return 0
 
+    @staticmethod
+    def gen_dict_extract(var, key):
+        '''
+        Generator function that yields all values for key in dictionary
+
+        :param var: dictionary or list of dictionaries
+
+        :param key: key name
+
+        :return: generator object containing values of key
+        '''
+        if isinstance(var, dict):
+            for k, v in var.items():
+                if k == key:
+                    yield v
+                if isinstance(v, (dict, list)):
+                    for result in PepperRetcode.gen_dict_extract(v, key):
+                        yield result
+        elif isinstance(var, list):
+            for d in var:
+                for result in PepperRetcode.gen_dict_extract(d, key):
+                    yield result
+
     @staticmethod
     def validate_fail_any(result):
         '''
@@ -40,14 +63,8 @@ def validate_fail_any(result):
 
         :return: exit code
         '''
-        if isinstance(result, list):
-            if isinstance(result[0], dict):
-                minion = result[0]
-                retcodes = list(minion[name].get('retcode')
-                                for name in minion if isinstance(minion[name], dict) and
-                                minion[name].get('retcode') is not None)
-                return next((r for r in retcodes if r != 0), 0)
-        return 0
+        retcodes = list(PepperRetcode.gen_dict_extract(result, 'retcode'))
+        return next((r for r in retcodes if r != 0), 0)
 
     @staticmethod
     def validate_fail_any_none(result):
@@ -60,16 +77,10 @@ def validate_fail_any_none(result):
 
         :return: exit code
         '''
-        if isinstance(result, list):
-            if isinstance(result[0], dict):
-                minion = result[0]
-                retcodes = list(minion[name].get('retcode')
-                                for name in minion if isinstance(minion[name], dict) and
-                                minion[name].get('retcode') is not None)
-                if not retcodes:
-                    return -1  # there are no retcodes
-                return next((r for r in retcodes if r != 0), 0)
-        return -1
+        retcodes = list(PepperRetcode.gen_dict_extract(result, 'retcode'))
+        if not retcodes:
+            return -1  # there are no retcodes
+        return next((r for r in retcodes if r != 0), 0)
 
     @staticmethod
     def validate_fail_all(result):
@@ -82,14 +93,9 @@ def validate_fail_all(result):
 
         :return: exit code
         '''
-        if isinstance(result, list):
-            if isinstance(result[0], dict):
-                minion = result[0]
-                retcodes = list(minion[name].get('retcode')
-                                for name in minion if isinstance(minion[name], dict) and
-                                minion[name].get('retcode') is not None)
-                if all(r != 0 for r in retcodes):
-                    return next((r for r in retcodes if r != 0), 0)
+        retcodes = list(PepperRetcode.gen_dict_extract(result, 'retcode'))
+        if all(r != 0 for r in retcodes):
+            return next((r for r in retcodes if r != 0), 0)
         return 0
 
     @staticmethod
@@ -103,16 +109,9 @@ def validate_fail_all_none(result):
 
         :return: exit code
         '''
-        if isinstance(result, list):
-            if isinstance(result[0], dict):
-                minion = result[0]
-                retcodes = list(minion[name].get('retcode')
-                                for name in minion if isinstance(minion[name], dict) and
-                                minion[name].get('retcode') is not None)
-                if not retcodes:
-                    return -1  # there are no retcodes
-                if all(r != 0 for r in retcodes):
-                    return next((r for r in retcodes if r != 0), 0)
-                else:
-                    return 0
-        return -1
+        retcodes = list(PepperRetcode.gen_dict_extract(result, 'retcode'))
+        if not retcodes:
+            return -1  # there are no retcodes
+        if all(r != 0 for r in retcodes):
+            return next((r for r in retcodes if r != 0), 0)
+        return 0