git.maemo.org Git - mevemon/blob - package/src/eveapi/eveapi.py

   1 #-----------------------------------------------------------------------------
   2 # eveapi - EVE Online API access
   3 #
   4 # Copyright (c)2007 Jamie "Entity" van den Berge <entity@vapor.com>
   5 #
   6 # Permission is hereby granted, free of charge, to any person
   7 # obtaining a copy of this software and associated documentation
   8 # files (the "Software"), to deal in the Software without
   9 # restriction, including without limitation the rights to use,
  10 # copy, modify, merge, publish, distribute, sublicense, and/or sell
  11 # copies of the Software, and to permit persons to whom the
  12 # Software is furnished to do so, subject to the following
  13 # conditions:
  14 #
  15 # The above copyright notice and this permission notice shall be
  16 # included in all copies or substantial portions of the Software.
  17 #
  18 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  19 # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
  20 # OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  21 # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
  22 # HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
  23 # WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  24 # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
  25 # OTHER DEALINGS IN THE SOFTWARE
  26 #
  27 #-----------------------------------------------------------------------------
  28 # Version: 1.1.9 - 2 September 2011
  29 # - added workaround for row tags with attributes that were not defined
  30 #   in their rowset (this should fix AssetList)
  31 #
  32 # Version: 1.1.8 - 1 September 2011
  33 # - fix for inconsistent columns attribute in rowsets.
  34 #
  35 # Version: 1.1.7 - 1 September 2011
  36 # - auth() method updated to work with the new authentication scheme.
  37 #
  38 # Version: 1.1.6 - 27 May 2011
  39 # - Now supports composite keys for IndexRowsets.
  40 # - Fixed calls not working if a path was specified in the root url.
  41 #
  42 # Version: 1.1.5 - 27 Januari 2011
  43 # - Now supports (and defaults to) HTTPS. Non-SSL proxies will still work by
  44 #   explicitly specifying http:// in the url.
  45 #
  46 # Version: 1.1.4 - 1 December 2010
  47 # - Empty explicit CDATA tags are now properly handled.
  48 # - _autocast now receives the name of the variable it's trying to typecast,
  49 #   enabling custom/future casting functions to make smarter decisions.
  50 #
  51 # Version: 1.1.3 - 6 November 2010
  52 # - Added support for anonymous CDATA inside row tags. This makes the body of
  53 #   mails in the rows of char/MailBodies available through the .data attribute.
  54 #
  55 # Version: 1.1.2 - 2 July 2010
  56 # - Fixed __str__ on row objects to work properly with unicode strings.
  57 #
  58 # Version: 1.1.1 - 10 Januari 2010
  59 # - Fixed bug that causes nested tags to not appear in rows of rowsets created
  60 #   from normal Elements. This should fix the corp.MemberSecurity method,
  61 #   which now returns all data for members. [jehed]
  62 #
  63 # Version: 1.1.0 - 15 Januari 2009
  64 # - Added Select() method to Rowset class. Using it avoids the creation of
  65 #   temporary row instances, speeding up iteration considerably.
  66 # - Added ParseXML() function, which can be passed arbitrary API XML file or
  67 #   string objects.
  68 # - Added support for proxy servers. A proxy can be specified globally or
  69 #   per api connection instance. [suggestion by graalman]
  70 # - Some minor refactoring.
  71 # - Fixed deprecation warning when using Python 2.6.
  72 #
  73 # Version: 1.0.7 - 14 November 2008
  74 # - Added workaround for rowsets that are missing the (required!) columns
  75 #   attribute. If missing, it will use the columns found in the first row.
  76 #   Note that this is will still break when expecting columns, if the rowset
  77 #   is empty. [Flux/Entity]
  78 #
  79 # Version: 1.0.6 - 18 July 2008
  80 # - Enabled expat text buffering to avoid content breaking up. [BigWhale]
  81 #
  82 # Version: 1.0.5 - 03 February 2008
  83 # - Added workaround to make broken XML responses (like the "row:name" bug in
  84 #   eve/CharacterID) work as intended.
  85 # - Bogus datestamps before the epoch in XML responses are now set to 0 to
  86 #   avoid breaking certain date/time functions. [Anathema Matou]
  87 #
  88 # Version: 1.0.4 - 23 December 2007
  89 # - Changed _autocast() to use timegm() instead of mktime(). [Invisible Hand]
  90 # - Fixed missing attributes of elements inside rows. [Elandra Tenari]
  91 #
  92 # Version: 1.0.3 - 13 December 2007
  93 # - Fixed keyless columns bugging out the parser (in CorporationSheet for ex.)
  94 #
  95 # Version: 1.0.2 - 12 December 2007
  96 # - Fixed parser not working with indented XML.
  97 #
  98 # Version: 1.0.1
  99 # - Some micro optimizations
 100 #
 101 # Version: 1.0
 102 # - Initial release
 103 #
 104 # Requirements:
 105 #   Python 2.4+
 106 #
 107 #-----------------------------------------------------------------------------
 108
 109 import httplib
 110 import urlparse
 111 import urllib
 112 import copy
 113
 114 from xml.parsers import expat
 115 from time import strptime
 116 from calendar import timegm
 117
 118 proxy = None
 119
 120 #-----------------------------------------------------------------------------
 121
 122 class Error(StandardError):
 123         def __init__(self, code, message):
 124                 self.code = code
 125                 self.args = (message.rstrip("."),)
 126
 127
 128 def EVEAPIConnection(url="api.eveonline.com", cacheHandler=None, proxy=None):
 129         # Creates an API object through which you can call remote functions.
 130         #
 131         # The following optional arguments may be provided:
 132         #
 133         # url - root location of the EVEAPI server
 134         #
 135         # proxy - (host,port) specifying a proxy server through which to request
 136         #         the API pages. Specifying a proxy overrides default proxy.
 137         #
 138         # cacheHandler - an object which must support the following interface:
 139         #
 140         #      retrieve(host, path, params)
 141         #
 142         #          Called when eveapi wants to fetch a document.
 143         #          host is the address of the server, path is the full path to
 144         #          the requested document, and params is a dict containing the
 145         #          parameters passed to this api call (keyID, vCode, etc).
 146         #          The method MUST return one of the following types:
 147         #
 148         #           None - if your cache did not contain this entry
 149         #           str/unicode - eveapi will parse this as XML
 150         #           Element - previously stored object as provided to store()
 151         #           file-like object - eveapi will read() XML from the stream.
 152         #
 153         #      store(host, path, params, doc, obj)
 154         #
 155         #          Called when eveapi wants you to cache this item.
 156         #          You can use obj to get the info about the object (cachedUntil
 157         #          and currentTime, etc) doc is the XML document the object
 158         #          was generated from. It's generally best to cache the XML, not
 159         #          the object, unless you pickle the object. Note that this method
 160         #          will only be called if you returned None in the retrieve() for
 161         #          this object.
 162         #
 163
 164         if not url.startswith("http"):
 165                 url = "https://" + url
 166         p = urlparse.urlparse(url, "https")
 167         if p.path and p.path[-1] == "/":
 168                 p.path = p.path[:-1]
 169         ctx = _RootContext(None, p.path, {}, {})
 170         ctx._handler = cacheHandler
 171         ctx._scheme = p.scheme
 172         ctx._host = p.netloc
 173         ctx._proxy = proxy or globals()["proxy"]
 174         return ctx
 175
 176
 177 def ParseXML(file_or_string):
 178         try:
 179                 return _ParseXML(file_or_string, False, None)
 180         except TypeError:
 181                 raise TypeError("XML data must be provided as string or file-like object")
 182
 183
 184 def _ParseXML(response, fromContext, storeFunc):
 185         # pre/post-process XML or Element data
 186
 187         if fromContext and isinstance(response, Element):
 188                 obj = response
 189         elif type(response) in (str, unicode):
 190                 obj = _Parser().Parse(response, False)
 191         elif hasattr(response, "read"):
 192                 obj = _Parser().Parse(response, True)
 193         else:
 194                 raise TypeError("retrieve method must return None, string, file-like object or an Element instance")
 195
 196         error = getattr(obj, "error", False)
 197         if error:
 198                 raise Error(error.code, error.data)
 199
 200         result = getattr(obj, "result", False)
 201         if not result:
 202                 raise RuntimeError("API object does not contain result")
 203
 204         if fromContext and storeFunc:
 205                 # call the cache handler to store this object
 206                 storeFunc(obj)
 207
 208         # make metadata available to caller somehow
 209         result._meta = obj
 210
 211         return result
 212
 213
 214
 215
 216
 217 #-----------------------------------------------------------------------------
 218 # API Classes
 219 #-----------------------------------------------------------------------------
 220
 221 _listtypes = (list, tuple, dict)
 222 _unspecified = []
 223
 224 class _Context(object):
 225
 226         def __init__(self, root, path, parentDict, newKeywords=None):
 227                 self._root = root or self
 228                 self._path = path
 229                 if newKeywords:
 230                         if parentDict:
 231                                 self.parameters = parentDict.copy()
 232                         else:
 233                                 self.parameters = {}
 234                         self.parameters.update(newKeywords)
 235                 else:
 236                         self.parameters = parentDict or {}
 237
 238         def context(self, *args, **kw):
 239                 if kw or args:
 240                         path = self._path
 241                         if args:
 242                                 path += "/" + "/".join(args)
 243                         return self.__class__(self._root, path, self.parameters, kw)
 244                 else:
 245                         return self
 246
 247         def __getattr__(self, this):
 248                 # perform arcane attribute majick trick
 249                 return _Context(self._root, self._path + "/" + this, self.parameters)
 250
 251         def __call__(self, **kw):
 252                 if kw:
 253                         # specified keywords override contextual ones
 254                         for k, v in self.parameters.iteritems():
 255                                 if k not in kw:
 256                                         kw[k] = v
 257                 else:
 258                         # no keywords provided, just update with contextual ones.
 259                         kw.update(self.parameters)
 260
 261                 # now let the root context handle it further
 262                 return self._root(self._path, **kw)
 263
 264
 265 class _AuthContext(_Context):
 266
 267         def character(self, characterID):
 268                 # returns a copy of this connection object but for every call made
 269                 # through it, it will add the folder "/char" to the url, and the
 270                 # characterID to the parameters passed.
 271                 return _Context(self._root, self._path + "/char", self.parameters, {"characterID":characterID})
 272
 273         def corporation(self, characterID):
 274                 # same as character except for the folder "/corp"
 275                 return _Context(self._root, self._path + "/corp", self.parameters, {"characterID":characterID})
 276
 277
 278 class _RootContext(_Context):
 279
 280         def auth(self, **kw):
 281                 if len(kw) == 2 and (("keyID" in kw and "vCode" in kw) or ("userID" in kw and "apiKey" in kw)):
 282                         return _AuthContext(self._root, self._path, self.parameters, kw)
 283                 raise ValueError("Must specify keyID and vCode")
 284
 285         def setcachehandler(self, handler):
 286                 self._root._handler = handler
 287
 288         def __call__(self, path, **kw):
 289                 # convert list type arguments to something the API likes
 290                 for k, v in kw.iteritems():
 291                         if isinstance(v, _listtypes):
 292                                 kw[k] = ','.join(map(str, list(v)))
 293
 294                 cache = self._root._handler
 295
 296                 # now send the request
 297                 path += ".xml.aspx"
 298
 299                 if cache:
 300                         response = cache.retrieve(self._host, path, kw)
 301                 else:
 302                         response = None
 303
 304                 if response is None:
 305                         if self._scheme == "https":
 306                                 connectionclass = httplib.HTTPSConnection
 307                         else:
 308                                 connectionclass = httplib.HTTPConnection
 309
 310                         if self._proxy is None:
 311                                 http = connectionclass(self._host)
 312                                 if kw:
 313                                         http.request("POST", path, urllib.urlencode(kw), {"Content-type": "application/x-www-form-urlencoded"})
 314                                 else:
 315                                         http.request("GET", path)
 316                         else:
 317                                 http = connectionclass(*self._proxy)
 318                                 if kw:
 319                                         http.request("POST", 'https://'+self._host+path, urllib.urlencode(kw), {"Content-type": "application/x-www-form-urlencoded"})
 320                                 else:
 321                                         http.request("GET", 'https://'+self._host+path)
 322
 323                         response = http.getresponse()
 324                         if response.status != 200:
 325                                 if response.status == httplib.NOT_FOUND:
 326                                         raise AttributeError("'%s' not available on API server (404 Not Found)" % path)
 327                                 else:
 328                                         raise RuntimeError("'%s' request failed (%d %s)" % (path, response.status, response.reason))
 329
 330                         if cache:
 331                                 store = True
 332                                 response = response.read()
 333                         else:
 334                                 store = False
 335                 else:
 336                         store = False
 337
 338                 retrieve_fallback = cache and getattr(cache, "retrieve_fallback", False)
 339                 if retrieve_fallback:
 340                         # implementor is handling fallbacks...
 341                         try:
 342                                 return _ParseXML(response, True, store and (lambda obj: cache.store(self._host, path, kw, response, obj)))
 343                         except Error, reason:
 344                                 response = retrieve_fallback(self._host, path, kw, reason=e)
 345                                 if response is not None:
 346                                         return response
 347                                 raise
 348                 else:
 349                         # implementor is not handling fallbacks...
 350                         return _ParseXML(response, True, store and (lambda obj: cache.store(self._host, path, kw, response, obj)))
 351
 352 #-----------------------------------------------------------------------------
 353 # XML Parser
 354 #-----------------------------------------------------------------------------
 355
 356 def _autocast(key, value):
 357         # attempts to cast an XML string to the most probable type.
 358         try:
 359                 if value.strip("-").isdigit():
 360                         return int(value)
 361         except ValueError:
 362                 pass
 363
 364         try:
 365                 return float(value)
 366         except ValueError:
 367                 pass
 368
 369         if len(value) == 19 and value[10] == ' ':
 370                 # it could be a date string
 371                 try:
 372                         return max(0, int(timegm(strptime(value, "%Y-%m-%d %H:%M:%S"))))
 373                 except OverflowError:
 374                         pass
 375                 except ValueError:
 376                         pass
 377
 378         # couldn't cast. return string unchanged.
 379         return value
 380
 381
 382
 383 class _Parser(object):
 384
 385         def Parse(self, data, isStream=False):
 386                 self.container = self.root = None
 387                 self._cdata = False
 388                 p = expat.ParserCreate()
 389                 p.StartElementHandler = self.tag_start
 390                 p.CharacterDataHandler = self.tag_cdata
 391                 p.StartCdataSectionHandler = self.tag_cdatasection_enter
 392                 p.EndCdataSectionHandler = self.tag_cdatasection_exit
 393                 p.EndElementHandler = self.tag_end
 394                 p.ordered_attributes = True
 395                 p.buffer_text = True
 396
 397                 if isStream:
 398                         p.ParseFile(data)
 399                 else:
 400                         p.Parse(data, True)
 401                 return self.root
 402
 403
 404         def tag_cdatasection_enter(self):
 405                 # encountered an explicit CDATA tag.
 406                 self._cdata = True
 407
 408         def tag_cdatasection_exit(self):
 409                 if self._cdata:
 410                         # explicit CDATA without actual data. expat doesn't seem
 411                         # to trigger an event for this case, so do it manually.
 412                         # (_cdata is set False by this call)
 413                         self.tag_cdata("")
 414                 else:
 415                         self._cdata = False
 416
 417         def tag_start(self, name, attributes):
 418                 # <hack>
 419                 # If there's a colon in the tag name, cut off the name from the colon
 420                 # onward. This is a workaround to make certain bugged XML responses
 421                 # (such as eve/CharacterID.xml.aspx) work.
 422                 if ":" in name:
 423                         name = name[:name.index(":")]
 424                 # </hack>
 425
 426                 if name == "rowset":
 427                         # for rowsets, use the given name
 428                         try:
 429                                 columns = attributes[attributes.index('columns')+1].replace(" ", "").split(",")
 430                         except ValueError:
 431                                 # rowset did not have columns tag set (this is a bug in API)
 432                                 # columns will be extracted from first row instead.
 433                                 columns = []
 434
 435                         try:
 436                                 priKey = attributes[attributes.index('key')+1]
 437                                 this = IndexRowset(cols=columns, key=priKey)
 438                         except ValueError:
 439                                 this = Rowset(cols=columns)
 440
 441
 442                         this._name = attributes[attributes.index('name')+1]
 443                         this.__catch = "row" # tag to auto-add to rowset.
 444                 else:
 445                         this = Element()
 446                         this._name = name
 447
 448                 this.__parent = self.container
 449
 450                 if self.root is None:
 451                         # We're at the root. The first tag has to be "eveapi" or we can't
 452                         # really assume the rest of the xml is going to be what we expect.
 453                         if name != "eveapi":
 454                                 raise RuntimeError("Invalid API response")
 455                         self.root = this
 456
 457                 if isinstance(self.container, Rowset) and (self.container.__catch == this._name):
 458                         # <hack>
 459                         # - check for missing columns attribute (see above)
 460                         # - check for extra attributes that were not defined in the rowset,
 461                         #   such as rawQuantity in the assets lists.
 462                         # In either case the tag is assumed to be correct and the rowset's
 463                         # columns are overwritten with the tag's version.
 464                         if not self.container._cols or (len(attributes)/2 > len(self.container._cols)):
 465                                 self.container._cols = attributes[0::2]
 466                         # </hack>
 467
 468                         self.container.append([_autocast(attributes[i], attributes[i+1]) for i in xrange(0, len(attributes), 2)])
 469                         this._isrow = True
 470                         this._attributes = this._attributes2 = None
 471                 else:
 472                         this._isrow = False
 473                         this._attributes = attributes
 474                         this._attributes2 = []
 475
 476                 self.container = this
 477
 478
 479         def tag_cdata(self, data):
 480                 if self._cdata:
 481                         # unset cdata flag to indicate it's been handled.
 482                         self._cdata = False
 483                 else:
 484                         if data in ("\r\n", "\n") or data.strip() != data:
 485                                 return
 486
 487                 this = self.container
 488                 data = _autocast(this._name, data)
 489
 490                 if this._isrow:
 491                         # sigh. anonymous data inside rows makes Entity cry.
 492                         # for the love of Jove, CCP, learn how to use rowsets.
 493                         parent = this.__parent
 494                         _row = parent._rows[-1]
 495                         _row.append(data)
 496                         if len(parent._cols) < len(_row):
 497                                 parent._cols.append("data")
 498
 499                 elif this._attributes:
 500                         # this tag has attributes, so we can't simply assign the cdata
 501                         # as an attribute to the parent tag, as we'll lose the current
 502                         # tag's attributes then. instead, we'll assign the data as
 503                         # attribute of this tag.
 504                         this.data = data
 505                 else:
 506                         # this was a simple <tag>data</tag> without attributes.
 507                         # we won't be doing anything with this actual tag so we can just
 508                         # bind it to its parent (done by __tag_end)
 509                         setattr(this.__parent, this._name, data)
 510
 511         def tag_end(self, name):
 512                 this = self.container
 513                 if this is self.root:
 514                         del this._attributes
 515                         #this.__dict__.pop("_attributes", None)
 516                         return
 517
 518                 # we're done with current tag, so we can pop it off. This means that
 519                 # self.container will now point to the container of element 'this'.
 520                 self.container = this.__parent
 521                 del this.__parent
 522
 523                 attributes = this.__dict__.pop("_attributes")
 524                 attributes2 = this.__dict__.pop("_attributes2")
 525                 if attributes is None:
 526                         # already processed this tag's closure early, in tag_start()
 527                         return
 528
 529                 if self.container._isrow:
 530                         # Special case here. tags inside a row! Such tags have to be
 531                         # added as attributes of the row.
 532                         parent = self.container.__parent
 533
 534                         # get the row line for this element from its parent rowset
 535                         _row = parent._rows[-1]
 536
 537                         # add this tag's value to the end of the row
 538                         _row.append(getattr(self.container, this._name, this))
 539
 540                         # fix columns if neccessary.
 541                         if len(parent._cols) < len(_row):
 542                                 parent._cols.append(this._name)
 543                 else:
 544                         # see if there's already an attribute with this name (this shouldn't
 545                         # really happen, but it doesn't hurt to handle this case!
 546                         sibling = getattr(self.container, this._name, None)
 547                         if sibling is None:
 548                                 self.container._attributes2.append(this._name)
 549                                 setattr(self.container, this._name, this)
 550                         # Note: there aren't supposed to be any NON-rowset tags containing
 551                         # multiples of some tag or attribute. Code below handles this case.
 552                         elif isinstance(sibling, Rowset):
 553                                 # its doppelganger is a rowset, append this as a row to that.
 554                                 row = [_autocast(attributes[i], attributes[i+1]) for i in xrange(0, len(attributes), 2)]
 555                                 row.extend([getattr(this, col) for col in attributes2])
 556                                 sibling.append(row)
 557                         elif isinstance(sibling, Element):
 558                                 # parent attribute is an element. This means we're dealing
 559                                 # with multiple of the same sub-tag. Change the attribute
 560                                 # into a Rowset, adding the sibling element and this one.
 561                                 rs = Rowset()
 562                                 rs.__catch = rs._name = this._name
 563                                 row = [_autocast(attributes[i], attributes[i+1]) for i in xrange(0, len(attributes), 2)]+[getattr(this, col) for col in attributes2]
 564                                 rs.append(row)
 565                                 row = [getattr(sibling, attributes[i]) for i in xrange(0, len(attributes), 2)]+[getattr(sibling, col) for col in attributes2]
 566                                 rs.append(row)
 567                                 rs._cols = [attributes[i] for i in xrange(0, len(attributes), 2)]+[col for col in attributes2]
 568                                 setattr(self.container, this._name, rs)
 569                         else:
 570                                 # something else must have set this attribute already.
 571                                 # (typically the <tag>data</tag> case in tag_data())
 572                                 pass
 573
 574                 # Now fix up the attributes and be done with it.
 575                 for i in xrange(0, len(attributes), 2):
 576                         this.__dict__[attributes[i]] = _autocast(attributes[i], attributes[i+1])
 577
 578                 return
 579
 580
 581
 582
 583 #-----------------------------------------------------------------------------
 584 # XML Data Containers
 585 #-----------------------------------------------------------------------------
 586 # The following classes are the various container types the XML data is
 587 # unpacked into.
 588 #
 589 # Note that objects returned by API calls are to be treated as read-only. This
 590 # is not enforced, but you have been warned.
 591 #-----------------------------------------------------------------------------
 592
 593 class Element(object):
 594         # Element is a namespace for attributes and nested tags
 595         def __str__(self):
 596                 return "<Element '%s'>" % self._name
 597
 598 _fmt = u"%s:%s".__mod__
 599 class Row(object):
 600         # A Row is a single database record associated with a Rowset.
 601         # The fields in the record are accessed as attributes by their respective
 602         # column name.
 603         #
 604         # To conserve resources, Row objects are only created on-demand. This is
 605         # typically done by Rowsets (e.g. when iterating over the rowset).
 606
 607         def __init__(self, cols=None, row=None):
 608                 self._cols = cols or []
 609                 self._row = row or []
 610
 611         def __nonzero__(self):
 612                 return True
 613
 614         def __ne__(self, other):
 615                 return self.__cmp__(other)
 616
 617         def __eq__(self, other):
 618                 return self.__cmp__(other) == 0
 619
 620         def __cmp__(self, other):
 621                 if type(other) != type(self):
 622                         raise TypeError("Incompatible comparison type")
 623                 return cmp(self._cols, other._cols) or cmp(self._row, other._row)
 624
 625         def __getattr__(self, this):
 626                 try:
 627                         return self._row[self._cols.index(this)]
 628                 except:
 629                         raise AttributeError, this
 630
 631         def __getitem__(self, this):
 632                 return self._row[self._cols.index(this)]
 633
 634         def __str__(self):
 635                 return "Row(" + ','.join(map(_fmt, zip(self._cols, self._row))) + ")"
 636
 637
 638 class Rowset(object):
 639         # Rowsets are collections of Row objects.
 640         #
 641         # Rowsets support most of the list interface:
 642         #   iteration, indexing and slicing
 643         #
 644         # As well as the following methods:
 645         #
 646         #   IndexedBy(column)
 647         #     Returns an IndexRowset keyed on given column. Requires the column to
 648         #     be usable as primary key.
 649         #
 650         #   GroupedBy(column)
 651         #     Returns a FilterRowset keyed on given column. FilterRowset objects
 652         #     can be accessed like dicts. See FilterRowset class below.
 653         #
 654         #   SortBy(column, reverse=True)
 655         #     Sorts rowset in-place on given column. for a descending sort,
 656         #     specify reversed=True.
 657         #
 658         #   SortedBy(column, reverse=True)
 659         #     Same as SortBy, except this returns a new rowset object instead of
 660         #     sorting in-place.
 661         #
 662         #   Select(columns, row=False)
 663         #     Yields a column values tuple (value, ...) for each row in the rowset.
 664         #     If only one column is requested, then just the column value is
 665         #     provided instead of the values tuple.
 666         #     When row=True, each result will be decorated with the entire row.
 667         #
 668
 669         def IndexedBy(self, column):
 670                 return IndexRowset(self._cols, self._rows, column)
 671
 672         def GroupedBy(self, column):
 673                 return FilterRowset(self._cols, self._rows, column)
 674
 675         def SortBy(self, column, reverse=False):
 676                 ix = self._cols.index(column)
 677                 self.sort(key=lambda e: e[ix], reverse=reverse)
 678
 679         def SortedBy(self, column, reverse=False):
 680                 rs = self[:]
 681                 rs.SortBy(column, reverse)
 682                 return rs
 683
 684         def Select(self, *columns, **options):
 685                 if len(columns) == 1:
 686                         i = self._cols.index(columns[0])
 687                         if options.get("row", False):
 688                                 for line in self._rows:
 689                                         yield (line, line[i])
 690                         else:
 691                                 for line in self._rows:
 692                                         yield line[i]
 693                 else:
 694                         i = map(self._cols.index, columns)
 695                         if options.get("row", False):
 696                                 for line in self._rows:
 697                                         yield line, [line[x] for x in i]
 698                         else:
 699                                 for line in self._rows:
 700                                         yield [line[x] for x in i]
 701
 702
 703         # -------------
 704
 705         def __init__(self, cols=None, rows=None):
 706                 self._cols = cols or []
 707                 self._rows = rows or []
 708
 709         def append(self, row):
 710                 if isinstance(row, list):
 711                         self._rows.append(row)
 712                 elif isinstance(row, Row) and len(row._cols) == len(self._cols):
 713                         self._rows.append(row._row)
 714                 else:
 715                         raise TypeError("incompatible row type")
 716
 717         def __add__(self, other):
 718                 if isinstance(other, Rowset):
 719                         if len(other._cols) == len(self._cols):
 720                                 self._rows += other._rows
 721                 raise TypeError("rowset instance expected")
 722
 723         def __nonzero__(self):
 724                 return not not self._rows
 725
 726         def __len__(self):
 727                 return len(self._rows)
 728
 729         def copy(self):
 730                 return self[:]
 731
 732         def __getitem__(self, ix):
 733                 if type(ix) is slice:
 734                         return Rowset(self._cols, self._rows[ix])
 735                 return Row(self._cols, self._rows[ix])
 736
 737         def sort(self, *args, **kw):
 738                 self._rows.sort(*args, **kw)
 739
 740         def __str__(self):
 741                 return ("Rowset(columns=[%s], rows=%d)" % (','.join(self._cols), len(self)))
 742
 743         def __getstate__(self):
 744                 return (self._cols, self._rows)
 745
 746         def __setstate__(self, state):
 747                 self._cols, self._rows = state
 748
 749
 750
 751 class IndexRowset(Rowset):
 752         # An IndexRowset is a Rowset that keeps an index on a column.
 753         #
 754         # The interface is the same as Rowset, but provides an additional method:
 755         #
 756         #   Get(key [, default])
 757         #     Returns the Row mapped to provided key in the index. If there is no
 758         #     such key in the index, KeyError is raised unless a default value was
 759         #     specified.
 760         #
 761
 762         def Get(self, key, *default):
 763                 row = self._items.get(key, None)
 764                 if row is None:
 765                         if default:
 766                                 return default[0]
 767                         raise KeyError, key
 768                 return Row(self._cols, row)
 769
 770         # -------------
 771
 772         def __init__(self, cols=None, rows=None, key=None):
 773                 try:
 774                         if "," in key:
 775                                 self._ki = ki = [cols.index(k) for k in key.split(",")]
 776                                 self.composite = True
 777                         else:
 778                                 self._ki = ki = cols.index(key)
 779                                 self.composite = False
 780                 except IndexError:
 781                         raise ValueError("Rowset has no column %s" % key)
 782
 783                 Rowset.__init__(self, cols, rows)
 784                 self._key = key
 785
 786                 if self.composite:
 787                         self._items = dict((tuple([row[k] for k in ki]), row) for row in self._rows)
 788                 else:
 789                         self._items = dict((row[ki], row) for row in self._rows)
 790
 791         def __getitem__(self, ix):
 792                 if type(ix) is slice:
 793                         return IndexRowset(self._cols, self._rows[ix], self._key)
 794                 return Rowset.__getitem__(self, ix)
 795
 796         def append(self, row):
 797                 Rowset.append(self, row)
 798                 if self.composite:
 799                         self._items[tuple([row[k] for k in self._ki])] = row
 800                 else:
 801                         self._items[row[self._ki]] = row
 802
 803         def __getstate__(self):
 804                 return (Rowset.__getstate__(self), self._items, self._ki)
 805
 806         def __setstate__(self, state):
 807                 state, self._items, self._ki = state
 808                 Rowset.__setstate__(self, state)
 809
 810
 811 class FilterRowset(object):
 812         # A FilterRowset works much like an IndexRowset, with the following
 813         # differences:
 814         # - FilterRowsets are accessed much like dicts
 815         # - Each key maps to a Rowset, containing only the rows where the value
 816         #   of the column this FilterRowset was made on matches the key.
 817
 818         def __init__(self, cols=None, rows=None, key=None, key2=None, dict=None):
 819                 if dict is not None:
 820                         self._items = items = dict
 821                 elif cols is not None:
 822                         self._items = items = {}
 823
 824                         idfield = cols.index(key)
 825                         if not key2:
 826                                 for row in rows:
 827                                         id = row[idfield]
 828                                         if id in items:
 829                                                 items[id].append(row)
 830                                         else:
 831                                                 items[id] = [row]
 832                         else:
 833                                 idfield2 = cols.index(key2)
 834                                 for row in rows:
 835                                         id = row[idfield]
 836                                         if id in items:
 837                                                 items[id][row[idfield2]] = row
 838                                         else:
 839                                                 items[id] = {row[idfield2]:row}
 840
 841                 self._cols = cols
 842                 self.key = key
 843                 self.key2 = key2
 844                 self._bind()
 845
 846         def _bind(self):
 847                 items = self._items
 848                 self.keys = items.keys
 849                 self.iterkeys = items.iterkeys
 850                 self.__contains__ = items.__contains__
 851                 self.has_key = items.has_key
 852                 self.__len__ = items.__len__
 853                 self.__iter__ = items.__iter__
 854
 855         def copy(self):
 856                 return FilterRowset(self._cols[:], None, self.key, self.key2, dict=copy.deepcopy(self._items))
 857
 858         def get(self, key, default=_unspecified):
 859                 try:
 860                         return self[key]
 861                 except KeyError:
 862                         if default is _unspecified:
 863                                 raise
 864                 return default
 865
 866         def __getitem__(self, i):
 867                 if self.key2:
 868                         return IndexRowset(self._cols, None, self.key2, self._items.get(i, {}))
 869                 return Rowset(self._cols, self._items[i])
 870
 871         def __getstate__(self):
 872                 return (self._cols, self._rows, self._items, self.key, self.key2)
 873
 874         def __setstate__(self, state):
 875                 self._cols, self._rows, self._items, self.key, self.key2 = state
 876                 self._bind()
 877