diff --git a/tools/visualize-pdr/pldm_visualise_pdrs.py b/tools/visualize-pdr/pldm_visualise_pdrs.py
index f81cf23..5ec8d88 100755
--- a/tools/visualize-pdr/pldm_visualise_pdrs.py
+++ b/tools/visualize-pdr/pldm_visualise_pdrs.py
@@ -126,6 +126,27 @@
                          view=False, cleanup=True, format='pdf')
 
 
+def get_pdrs(client):
+    """ Using pldmtool over SSH, generate (record handle, PDR) tuples for each
+        record in the PDR repository.
+
+        Parameters:
+            client: paramiko ssh client object
+
+    """
+
+    command_fmt = 'pldmtool platform getpdr -d {}'
+    record_handle = 0
+    while True:
+        output = client.exec_command(command_fmt.format(str(record_handle)))
+        _, stdout, stderr = output
+        pdr = json.load(stdout)
+        yield record_handle, pdr
+        record_handle = pdr["nextRecordHandle"]
+        if record_handle == 0:
+            break
+
+
 def fetch_pdrs_from_bmc(client):
 
     """ This is the core function that would use the existing ssh connection
@@ -145,14 +166,7 @@
     numeric_pdr = {}
     fru_record_set_pdr = {}
     tl_pdr = {}
-    handle_number = 0
-    while True:
-        my_str = ''
-        command = 'pldmtool platform getpdr -d ' + str(handle_number)
-        output = client.exec_command(command)
-        for line in output[1]:
-            my_str += line.strip('\n')
-        my_dic = json.loads(my_str)
+    for handle_number, my_dic in get_pdrs(client):
         sys.stdout.write("Fetching PDR's from BMC : %8d\r" % (handle_number))
         sys.stdout.flush()
         if my_dic["PDRType"] == "Entity Association PDR":
@@ -167,10 +181,6 @@
             tl_pdr[handle_number] = my_dic
         if my_dic["PDRType"] == "Numeric Effecter PDR":
             numeric_pdr[handle_number] = my_dic
-        if not my_dic["nextRecordHandle"] == 0:
-            handle_number = my_dic["nextRecordHandle"]
-        else:
-            break
     client.close()
 
     total_pdrs = len(entity_association_pdr.keys()) + len(tl_pdr.keys()) + \