- fix add_object() usage

[plcapi.git] / PLC / Methods / GetPCUs.py

diff --git a/PLC/Methods/GetPCUs.py b/PLC/Methods/GetPCUs.py
index 6da58a0..13bd7dc 100644 (file)
--- a/PLC/Methods/GetPCUs.py
+++ b/PLC/Methods/GetPCUs.py
@@ -1,13 +1,17 @@
 from PLC.Faults import *
 from PLC.Method import Method
 from PLC.Parameter import Parameter, Mixed
+from PLC.Filter import Filter
 from PLC.PCUs import PCU, PCUs
-from PLC.Auth import PasswordAuth
+from PLC.Auth import Auth
 
 class GetPCUs(Method):
     """
-    Return an array of structs containing details about PCUs. If
-    pcu_id_list is specified, only the specified PCUs will be queried.
+    Returns an array of structs containing details about power control
+    units (PCUs). If pcu_filter is specified and is an array of PCU
+    identifiers, or a struct of PCU attributes, only PCUs matching the
+    filter will be returned. If return_fields is specified, only the
+    specified details will be returned.
 
     Admin may query all PCUs. Non-admins may only query the PCUs at
     their sites.
@@ -16,29 +20,34 @@ class GetPCUs(Method):
     roles = ['admin', 'pi', 'tech']
 
     accepts = [
-        PasswordAuth(),
-        [PCU.fields['pcu_id']]
+        Auth(),
+        Mixed([PCU.fields['pcu_id']],
+              Filter(PCU.fields)),
+        Parameter([str], "List of fields to return", nullok = True)
         ]
 
     returns = [PCU.fields]
 
-    def call(self, auth, pcu_ids = None):
+    def call(self, auth, pcu_filter = None, return_fields = None):
        # If we are not admin, make sure to only return our own PCUs
         if 'admin' not in self.caller['roles']:
             # Get list of PCUs that we are able to view
-            sites = Sites(self.api, self.caller['site_ids']).values()
+            valid_pcu_ids = []
+            if self.caller['site_ids']:
+                sites = Sites(self.api, self.caller['site_ids'])
+                for site in sites:
+                    valid_pcu_ids += site['pcu_ids']
 
-            valid_pcu_ids = set()
-            for site in sites:
-                valid_pcu_ids = valid_pcu_ids.union(site['pcu_ids'])
+            if not valid_pcu_ids:
+                return []
 
-            if not pcu_ids:
-                pcu_ids = valid_pcu_ids
-            else:
-                pcu_ids = valid_pcu_ids.intersection(pcu_ids)
+            if pcu_filter is None:
+                pcu_filter = valid_pcu_ids
 
-        pcus = PCUs(self.api, pcu_ids).values()
+        pcus = PCUs(self.api, pcu_filter, return_fields)
 
-       pcus = [dict(pcu) for pcu in pcus]
+        # Filter out PCUs that are not viewable
+        if 'admin' not in self.caller['roles']:
+            pcus = filter(lambda pcu: pcu['pcu_id'] in valid_pcu_ids, pcus)
 
-       return pcus
+        return pcus