Improvements to the TradeCalc API

CHANGES.txt

@@ -4,6 +4,7 @@ TradeDangerous, Copyright (C) Oliver "kfsone" Smith, July 2014
 
 [wip]
 . (kfsone) Fixed 'set changed size' error in run command
+. (kfsone) Improvements to the TradeCalc api for accessing station trades
 
 v6.9.2 Feb 08 2015
 . (kfsone) Feature #158 Find rares based on distance from multiple systems

tradecalc.py

@@ -360,7 +360,26 @@ class TradeCalc(object):
     Container for accessing trade calculations with common properties.
     """
 
-    def __init__(self, tdb, tdenv, fit=None):
+    def __init__(self, tdb, tdenv, fit=None, items=None):
+        """
+        Constructs the TradeCalc object and loads sell/buy data.
+
+        Parameters:
+            tdb
+                The TradeDB() object to use to access data,
+            tdenv
+                TradeEnv() that controls behavior,
+            fit [optional]
+                Lets you specify a fitting function,
+            items [optional]
+                Iterable [itemID or Item()] that restricts loading,
+
+        TradeEnv options:
+            tdenv.avoidItems
+                Iterable of [Item] that prevents items being loaded
+            tdenv.maxAge
+                Maximum age in days of data that gets loaded
+        """
         self.tdb = tdb
         self.tdenv = tdenv
         self.defaultFit = fit or self.fastFit
@@ -372,11 +391,21 @@ def __init__(self, tdb, tdenv, fit=None):
         avoidItemIDs = set([item.ID for item in tdenv.avoidItems])
 
         if tdenv.maxAge:
-            loadWhere = "WHERE JULIANDAY(modified) >= JULIANDAY('NOW') - {:f}".format(
+            ageClause = "AND JULIANDAY(modified) >= JULIANDAY('NOW') - {:f}".format(
                     tdenv.maxAge
             )
         else:
-            loadWhere = ""
+            ageClause = ""
+
+        loadItems = items or tdb.itemByID.values()
+        loadItemIDs = set()
+        for item in loadItems:
+            ID = item if isinstance(item, int) else item.ID
+            if ID not in avoidItemIDs:
+                loadItemIDs.add(str(ID))
+        if not loadItemIDs:
+            raise TradeException("No items to load.")
+        loadItemIDs = ",".join(str(ID) for ID in loadItemIDs)
 
         def load_items(tableName, index):
             lastStnID, stnAppend = 0, None
@@ -387,23 +416,28 @@ def load_items(tableName, index):
                             strftime('%s', modified),
                             modified
                       FROM  {}
-                      {where}
-            """.format(tableName, where=loadWhere))
+                     WHERE  item_id IN ({ids})
+                      {ageClause}
+            """.format(
+                tableName,
+                ageClause=ageClause,
+                ids=loadItemIDs,
+                )
+            )
             now = int(time.time())
             for stnID, itmID, cr, units, lev, timestamp, modified in cur:
-                if itmID not in avoidItemIDs:
-                    if stnID != lastStnID:
-                        stnAppend = index[stnID].append
-                        lastStnID = stnID
-                    try:
-                        ageS = now - int(timestamp)
-                    except TypeError:
-                        raise BadTimestampError(
-                            TableName, self.tdb,
-                            stnID, itmID, modified
-                        )
-                    stnAppend((itmID, cr, units, lev, ageS))
-                    count += 1
+                if stnID != lastStnID:
+                    stnAppend = index[stnID].append
+                    lastStnID = stnID
+                try:
+                    ageS = now - int(timestamp)
+                except TypeError:
+                    raise BadTimestampError(
+                        TableName, self.tdb,
+                        stnID, itmID, modified
+                    )
+                stnAppend((itmID, cr, units, lev, ageS))
+                count += 1
             tdenv.DEBUG0("Loaded {} selling values".format(count))
 
         self.stationsSelling = defaultdict(list)