- use PLC_WWW_SSL_PORT

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

diff --git a/PLC/Methods/GetPCUs.py b/PLC/Methods/GetPCUs.py
index 6c8731b..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,24 +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
             valid_pcu_ids = []
             if self.caller['site_ids']:
-                sites = Sites(self.api, self.caller['site_ids']).values()
+                sites = Sites(self.api, self.caller['site_ids'])
                 for site in sites:
                     valid_pcu_ids += site['pcu_ids']
 
-            pcu_ids = set(pcu_ids).intersection(valid_pcu_ids)
-            if not pcu_ids:
+            if not valid_pcu_ids:
                 return []
 
-        return PCUs(self.api, pcu_ids).values()
+            if pcu_filter is None:
+                pcu_filter = valid_pcu_ids
+
+        pcus = PCUs(self.api, pcu_filter, return_fields)
+
+        # 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