website/archive/binaries/mac/src/disembl/biopython-1.50/Bio/triemodule.c

   1 #include <Python.h>
   2 #include <marshal.h>
   3 #include "trie.h"
   4
   5 #if PY_VERSION_HEX < 0x02050000
   6 #define Py_ssize_t int
   7 #endif
   8
   9
  10
  11 staticforward PyTypeObject Trie_Type;
  12
  13 typedef struct {
  14     PyObject_HEAD
  15     Trie trie;
  16 } trieobject;
  17
  18 static PyObject*
  19 trie_trie(PyObject* self, PyObject* args)
  20 {
  21     trieobject* trieobj;
  22     Trie trie;
  23
  24     if (!PyArg_ParseTuple(args,":trie"))
  25         return NULL;
  26     if(!(trie = Trie_new()))
  27         return PyErr_NoMemory();
  28     if(!(trieobj = PyObject_New(trieobject, &Trie_Type)))
  29         return NULL;
  30     trieobj->trie = trie;
  31     return (PyObject*)trieobj;
  32 }
  33
  34 static void
  35 _decref_objects(const unsigned char *key, const void *value, void *data)
  36 {
  37     Py_DECREF((PyObject *)value);
  38 }
  39
  40 static void
  41 trie_dealloc(PyObject* self)
  42 {
  43     trieobject *mp = (trieobject *)self;
  44     Trie_iterate(mp->trie, _decref_objects, NULL);
  45     Trie_del(mp->trie);
  46     PyObject_Del(self);
  47 }
  48
  49 static Py_ssize_t
  50 trie_length(trieobject *mp)
  51 {
  52     return Trie_len(mp->trie);
  53 }
  54
  55 static PyObject *
  56 trie_subscript(trieobject *mp, PyObject *py_key)
  57 {
  58     unsigned char *key;
  59     PyObject *py_value;
  60
  61     /* Make sure key is a string. */
  62     if(!PyString_Check(py_key)) {
  63         PyErr_SetString(PyExc_TypeError, "key must be a string");
  64         return NULL;
  65     }
  66     key = (unsigned char *)PyString_AS_STRING(py_key);
  67     py_value = (PyObject *)Trie_get(mp->trie, key);
  68     if(py_value == NULL)
  69         PyErr_SetString(PyExc_KeyError, (char *)key);
  70     else
  71         Py_INCREF(py_value);
  72     return py_value;
  73 }
  74
  75 static int
  76 trie_ass_sub(trieobject *mp, PyObject *py_key, PyObject *py_value)
  77 {
  78     unsigned char *key;
  79     PyObject *py_prev;
  80
  81     /* Make sure key is a string. */
  82     if(!PyString_Check(py_key)) {
  83         PyErr_SetString(PyExc_TypeError, "key must be a string");
  84         return -1;
  85     }
  86     key = (unsigned char *)PyString_AS_STRING((char *)py_key);
  87
  88     /* Check to see whether something already exists at that key.  If
  89        there's already an object there, then I will have to remove it.
  90     */
  91     py_prev = (PyObject *)Trie_get(mp->trie, key);
  92     if(py_prev) {
  93         Py_DECREF(py_prev);
  94     }
  95
  96     /* The client wants to delete a key from a dictionary.  The Trie
  97        API doesn't support this, so I will just overwrite it with
  98        NULL. */
  99     if(!py_value) {
 100         /* If the key doesn't exist, raise a KeyError. */
 101         if(!py_prev) {
 102             PyErr_SetString(PyExc_KeyError, (char *)key);
 103             return -1;
 104         }
 105         Trie_set(mp->trie, key, NULL);
 106     }
 107     /* The client wants to set a key in the dictionary. */
 108     else {
 109         Py_INCREF(py_value);
 110         if(Trie_set(mp->trie, key, py_value)) {
 111             PyErr_SetString(PyExc_AssertionError, "error setting trie");
 112             return -1;
 113         }
 114     }
 115     return 0;
 116 }
 117
 118 static char has_key__doc__[] =
 119 "D.has_key(k) -> 1 if D has a key k, else 0";
 120
 121 static PyObject *
 122 trie_has_key(trieobject *mp, PyObject *py_key)
 123 {
 124     unsigned char *key;
 125     int has_key;
 126
 127     /* Make sure key is a string. */
 128     if(!PyString_Check(py_key)) {
 129         PyErr_SetString(PyExc_TypeError, "key must be a string");
 130         return NULL;
 131     }
 132     key = (unsigned char *)PyString_AS_STRING(py_key);
 133     has_key = Trie_has_key(mp->trie, key);
 134     return PyInt_FromLong((long)has_key);
 135 }
 136
 137 static PyObject *
 138 trie_has_key_onearg(trieobject *mp, PyObject *py_args)
 139 {
 140     PyObject *py_arg;
 141     if(!PyArg_ParseTuple(py_args, "O", &py_arg))
 142         return NULL;
 143     return trie_has_key(mp, py_arg);
 144 }
 145
 146
 147
 148 static char has_prefix__doc__[] =
 149 "D.has_prefix(k) -> 1 if D has a prefix k, else 0";
 150
 151 static PyObject *
 152 trie_has_prefix(trieobject *mp, PyObject *py_prefix)
 153 {
 154     unsigned char *prefix;
 155     int has_prefix;
 156
 157     /* Make sure prefix is a string. */
 158     if(!PyString_Check(py_prefix)) {
 159         PyErr_SetString(PyExc_TypeError, "k must be a string");
 160         return NULL;
 161     }
 162     prefix = (unsigned char *)PyString_AS_STRING(py_prefix);
 163     has_prefix = Trie_has_prefix(mp->trie, prefix);
 164     return PyInt_FromLong((long)has_prefix);
 165 }
 166
 167 static PyObject *
 168 trie_has_prefix_onearg(trieobject *mp, PyObject *py_args)
 169 {
 170     PyObject *py_arg;
 171     if(!PyArg_ParseTuple(py_args, "O", &py_arg))
 172         return NULL;
 173     return trie_has_prefix(mp, py_arg);
 174 }
 175
 176 static char with_prefix__doc__[] =
 177 "D.with_prefix(prefix) -> list of D's keys that begins with prefix";
 178
 179 static void
 180 _trie_with_prefix_helper(const unsigned char *key, const void *value,
 181                          void *data)
 182 {
 183     PyObject *py_list = (PyObject *)data;
 184     PyObject *py_key;
 185
 186     if(PyErr_Occurred())
 187         return;
 188
 189     if(!(py_key = PyString_FromString((const char *)key)))
 190         return;
 191     PyList_Append(py_list, py_key);
 192     Py_DECREF(py_key);
 193 }
 194
 195 static PyObject *
 196 trie_with_prefix(trieobject *mp, PyObject *py_prefix)
 197 {
 198     unsigned char *prefix;
 199     PyObject *py_list;
 200
 201     /* Make sure prefix is a string. */
 202     if(!PyString_Check(py_prefix)) {
 203         PyErr_SetString(PyExc_TypeError, "k must be a string");
 204         return NULL;
 205     }
 206     prefix = (unsigned char *)PyString_AS_STRING(py_prefix);
 207
 208     if(!(py_list = PyList_New(0)))
 209         return NULL;
 210     Trie_with_prefix(mp->trie, prefix,
 211                      _trie_with_prefix_helper, (void *)py_list);
 212     if(PyErr_Occurred()) {
 213         Py_DECREF(py_list);
 214         return NULL;
 215     }
 216     return py_list;
 217 }
 218
 219 static PyObject *
 220 trie_with_prefix_onearg(trieobject *mp, PyObject *py_args)
 221 {
 222     PyObject *py_arg;
 223     if(!PyArg_ParseTuple(py_args, "O", &py_arg))
 224         return NULL;
 225     return trie_with_prefix(mp, py_arg);
 226 }
 227
 228
 229 static char keys__doc__[] =
 230 "D.keys() -> list of D's keys";
 231
 232 static void
 233 _trie_keys_helper(const unsigned char *key, const void *value, void *data)
 234 {
 235     PyObject *py_list = (PyObject *)data;
 236     PyObject *py_key;
 237
 238     if(PyErr_Occurred())
 239         return;
 240
 241     if(!(py_key = PyString_FromString((char *)key)))
 242         return;
 243     PyList_Append(py_list, py_key);
 244     Py_DECREF(py_key);
 245 }
 246
 247 static PyObject *
 248 trie_keys(trieobject *mp)
 249 {
 250     PyObject *py_list;
 251
 252     if(!(py_list = PyList_New(0)))
 253         return NULL;
 254     Trie_iterate(mp->trie, _trie_keys_helper, (void *)py_list);
 255     if(PyErr_Occurred()) {
 256         Py_DECREF(py_list);
 257         return NULL;
 258     }
 259     return py_list;
 260 }
 261
 262 static PyObject *
 263 trie_keys_noargs(trieobject *mp, PyObject *py_args)
 264 {
 265     if(PyTuple_Size(py_args) != 0) {
 266         PyErr_SetString(PyExc_ValueError, "no args expected");
 267         return NULL;
 268     }
 269     return trie_keys(mp);
 270 }
 271
 272 static char values__doc__[] =
 273 "D.values() -> list of D's values";
 274
 275 static void
 276 _trie_values_helper(const unsigned char *key, const void *value, void *data)
 277 {
 278     PyObject *py_list = (PyObject *)data;
 279     if(PyErr_Occurred())
 280         return;
 281     PyList_Append(py_list, (PyObject *)value);
 282 }
 283
 284 static PyObject *
 285 trie_values(trieobject *mp)
 286 {
 287     PyObject *py_list;
 288
 289     if(!(py_list = PyList_New(0)))
 290         return NULL;
 291     Trie_iterate(mp->trie, _trie_values_helper, (void *)py_list);
 292     if(PyErr_Occurred()) {
 293         Py_DECREF(py_list);
 294         return NULL;
 295     }
 296     return py_list;
 297 }
 298
 299 static PyObject *
 300 trie_values_noargs(trieobject *mp, PyObject *py_args)
 301 {
 302     if(PyTuple_Size(py_args) != 0) {
 303         PyErr_SetString(PyExc_ValueError, "no args expected");
 304         return NULL;
 305     }
 306     return trie_values(mp);
 307 }
 308
 309 static char get__doc__[] =
 310 "D.get(k[,d]) -> D[k] if D.has_key(k), else d.  d defaults to None.";
 311
 312 static PyObject *
 313 trie_get(trieobject *mp, PyObject *args)
 314 {
 315     unsigned char *key;
 316     PyObject *py_value;
 317     PyObject *py_failobj = Py_None;
 318
 319     if (!PyArg_ParseTuple(args, "s|O:get", &key, &py_failobj))
 320         return NULL;
 321     py_value = (PyObject *)Trie_get(mp->trie, key);
 322     if(!py_value)
 323         py_value = py_failobj;
 324     Py_INCREF(py_value);
 325     return py_value;
 326 }
 327
 328 static char get_approximate__doc__[] =
 329 "D.get_approximate(key, k) -> List of (key, value, mismatches) in D, allowing up to k mismatches in key.";
 330
 331 void
 332 _trie_get_approximate_helper(const unsigned char *key, const void *value,
 333                              const int mismatches, void *data)
 334 {
 335     /* Append a tuple of (key, value) to data, which is a PyList. */
 336     PyObject *py_list = (PyObject *)data,
 337         *py_value = (PyObject *)value,
 338         *py_key,
 339         *py_tuple,
 340         *py_mismatches;
 341
 342     if(PyErr_Occurred())
 343         return;
 344
 345     if(!(py_key = PyString_FromString((const char *)key)))
 346         return;
 347     if(!(py_mismatches = PyInt_FromLong(mismatches))) {
 348         Py_DECREF(py_key);
 349         return;
 350     }
 351     Py_INCREF(py_value);
 352
 353     if(!(py_tuple = PyTuple_New(3))) {
 354         Py_DECREF(py_key);
 355         Py_DECREF(py_mismatches);
 356         Py_DECREF(py_value);
 357         return;
 358     }
 359     PyTuple_SetItem(py_tuple, 0, py_key);
 360     PyTuple_SetItem(py_tuple, 1, py_value);
 361     PyTuple_SetItem(py_tuple, 2, py_mismatches);
 362     PyList_Append(py_list, py_tuple);
 363     Py_DECREF(py_tuple);
 364 }
 365
 366 static PyObject *
 367 trie_get_approximate(trieobject *mp, PyObject *args)
 368 {
 369     unsigned char *key;
 370     int k;
 371     PyObject *py_list;
 372
 373     if (!PyArg_ParseTuple(args, "si:get_approximate", &key, &k))
 374         return NULL;
 375
 376     if(!(py_list = PyList_New(0)))
 377         return NULL;
 378     Trie_get_approximate(mp->trie, key, k,
 379                          _trie_get_approximate_helper, (void *)py_list);
 380     if(PyErr_Occurred()) {
 381         Py_DECREF(py_list);
 382         return NULL;
 383     }
 384     return py_list;
 385 }
 386
 387 static long
 388 trie_nohash(PyObject *self)
 389 {
 390     PyErr_SetString(PyExc_TypeError, "trie objects are unhashable");
 391     return -1;
 392 }
 393
 394 static PyMappingMethods trie_as_mapping = {
 395 /* The first member of PyMappingMethods was redefined in Python 2.5. */
 396 #if PY_VERSION_HEX < 0x02050000
 397     (inquiry)trie_length,        /*mp_length*/
 398 #else
 399     (lenfunc)trie_length,        /*mp_length*/
 400 #endif
 401     (binaryfunc)trie_subscript,  /*mp_subscript*/
 402     (objobjargproc)trie_ass_sub  /*mp_ass_subscript*/
 403 };
 404
 405 static PyMethodDef trieobj_methods[] = {
 406     /*  METH_O and METH_NOARGS require Python 2.2.
 407     {"has_key", (PyCFunction)trie_has_key,  METH_O,
 408      has_key__doc__},
 409     {"has_prefix", (PyCFunction)trie_has_prefix,  METH_O,
 410      has_prefix__doc__},
 411     {"with_prefix", (PyCFunction)trie_with_prefix,  METH_O,
 412      with_prefix__doc__},
 413     {"keys",    (PyCFunction)trie_keys,     METH_NOARGS,
 414      keys__doc__},
 415     {"values",  (PyCFunction)trie_values,   METH_NOARGS,
 416      values__doc__},
 417     */
 418
 419     {"has_key", (PyCFunction)trie_has_key_onearg,  METH_VARARGS,
 420      has_key__doc__},
 421     {"has_prefix", (PyCFunction)trie_has_prefix_onearg,  METH_VARARGS,
 422      has_prefix__doc__},
 423     {"with_prefix", (PyCFunction)trie_with_prefix_onearg,  METH_VARARGS,
 424      with_prefix__doc__},
 425     {"keys",    (PyCFunction)trie_keys_noargs,     METH_VARARGS,
 426      keys__doc__},
 427     {"values",  (PyCFunction)trie_values_noargs,   METH_VARARGS,
 428      values__doc__},
 429
 430     {"get",     (PyCFunction)trie_get,      METH_VARARGS,
 431      get__doc__},
 432     {"get_approximate",  (PyCFunction)trie_get_approximate,  METH_VARARGS,
 433      get_approximate__doc__},
 434     {NULL, NULL}   /* sentinel */
 435 };
 436
 437 static PyObject *trie_getattr(PyObject *obj, char *name)
 438 {
 439     return Py_FindMethod(trieobj_methods, (PyObject *)obj, name);
 440
 441 }
 442
 443 static PyTypeObject Trie_Type = {
 444     PyObject_HEAD_INIT(NULL)
 445     0,
 446     "trie",
 447     sizeof(trieobject),
 448     0,
 449     trie_dealloc,       /*tp_dealloc*/
 450     0,                  /*tp_print*/
 451     trie_getattr,                  /*tp_getattr*/
 452     0,                  /*tp_setattr*/
 453     0,                  /*tp_compare*/
 454     0,                  /*tp_repr*/
 455     0,                  /*tp_as_number*/
 456     0,                  /*tp_as_sequence*/
 457     &trie_as_mapping,   /*tp_as_mapping*/
 458     trie_nohash,        /*tp_hash */
 459 };
 460
 461 static int
 462 _write_to_handle(const void *towrite, const int length, void *handle)
 463 {
 464     PyObject *py_handle = (PyObject *)handle,
 465         *py_retval = NULL;
 466     int success = 0;
 467
 468     if(!length)
 469         return 1;
 470
 471     if(!(py_retval = PyObject_CallMethod(py_handle, "write", "s#",
 472                                          towrite, length)))
 473         goto _write_to_handle_cleanup;
 474     success = 1;
 475
 476  _write_to_handle_cleanup:
 477     if(py_retval) {
 478         Py_DECREF(py_retval);
 479     }
 480     return success;
 481 }
 482
 483 int _write_value_to_handle(const void *value, void *handle)
 484 {
 485     PyObject *py_value = (PyObject *)value,
 486         *py_marshalled = NULL;
 487     char *marshalled;
 488     Py_ssize_t length;
 489     int success = 0;
 490
 491 #ifdef Py_MARSHAL_VERSION
 492     if(!(py_marshalled =
 493          PyMarshal_WriteObjectToString(py_value, Py_MARSHAL_VERSION)))
 494         goto _write_value_to_handle_cleanup;
 495 #else
 496     if(!(py_marshalled = PyMarshal_WriteObjectToString(py_value)))
 497         goto _write_value_to_handle_cleanup;
 498 #endif
 499     if(PyString_AsStringAndSize(py_marshalled, &marshalled, &length) == -1)
 500         goto _write_value_to_handle_cleanup;
 501     if(!_write_to_handle(&length, sizeof(length), handle))
 502         goto _write_value_to_handle_cleanup;
 503     if (length != (int)length)
 504         goto _write_value_to_handle_cleanup;
 505     if(!_write_to_handle(marshalled, (int)length, handle))
 506         goto _write_value_to_handle_cleanup;
 507     success = 1;
 508
 509  _write_value_to_handle_cleanup:
 510     if(py_marshalled) {
 511         Py_DECREF(py_marshalled);
 512     }
 513
 514     return success;
 515 }
 516
 517 static PyObject *
 518 trie_save(PyObject *self, PyObject *args)
 519 {
 520     PyObject *py_handle,
 521         *py_trie;
 522     trieobject *mp;
 523
 524     if(!PyArg_ParseTuple(args, "OO:save", &py_handle, &py_trie))
 525         return NULL;
 526     mp = (trieobject *)py_trie;
 527     if(!Trie_serialize(mp->trie, _write_to_handle, _write_value_to_handle,
 528                        (void *)py_handle)) {
 529         if(!PyErr_Occurred())
 530             PyErr_SetString(PyExc_RuntimeError,
 531                             "saving failed for some reason");
 532         return NULL;
 533     }
 534     Py_INCREF(Py_None);
 535     return Py_None;
 536 }
 537
 538 static int
 539 _read_from_handle(void *wasread, const int length, void *handle)
 540 {
 541     PyObject *py_handle = (PyObject *)handle,
 542         *py_retval = NULL;
 543     void *retval;
 544     int success = 0;
 545     PyBufferProcs *buffer;
 546     int segment;
 547     int bytes_read, bytes_left;
 548
 549     if(!length)
 550         return 1;
 551
 552     if(!(py_retval = PyObject_CallMethod(py_handle, "read", "i", length)))
 553         goto _read_from_handle_cleanup;
 554     if(!py_retval->ob_type->tp_as_buffer) {
 555         PyErr_SetString(PyExc_ValueError, "read method should return buffer");
 556         goto _read_from_handle_cleanup;
 557     }
 558     if(!(py_retval->ob_type->tp_flags & Py_TPFLAGS_DEFAULT)) {
 559         PyErr_SetString(PyExc_ValueError, "no bf_getcharbuffer slot");
 560         goto _read_from_handle_cleanup;
 561     }
 562     buffer = py_retval->ob_type->tp_as_buffer;
 563     if(!buffer->bf_getreadbuffer) {
 564         PyErr_SetString(PyExc_ValueError, "no bf_getreadbuffer");
 565         goto _read_from_handle_cleanup;
 566     }
 567
 568     bytes_left = length;
 569     segment = 0;
 570     while(bytes_left > 0) {
 571         if((bytes_read = buffer->bf_getreadbuffer(py_retval,
 572                                                   segment, &retval)) == -1)
 573             goto _read_from_handle_cleanup;
 574         memcpy(wasread, retval, bytes_read);
 575         wasread = (void *)((char *)wasread + bytes_read);
 576         bytes_left -= bytes_read;
 577         segment += 1;
 578     }
 579
 580     success = 1;
 581
 582  _read_from_handle_cleanup:
 583     if(py_retval) {
 584         Py_DECREF(py_retval);
 585     }
 586     return success;
 587 }
 588
 589 #define MAX_KEY_LENGTH 2000
 590 static void *
 591 _read_value_from_handle(void *handle)
 592 {
 593     Py_ssize_t length;
 594     char KEY[MAX_KEY_LENGTH];
 595
 596     if(!_read_from_handle((void *)&length, sizeof(length), (void *)handle))
 597         return NULL;
 598     if(length < 0 || length >= MAX_KEY_LENGTH)
 599         return NULL;
 600     if(!_read_from_handle((void *)KEY, length, (void *)handle))
 601         return NULL;
 602     return PyMarshal_ReadObjectFromString(KEY, length);
 603 }
 604
 605
 606 static PyObject *
 607 trie_load(PyObject *self, PyObject *args)
 608 {
 609     PyObject *py_handle;
 610     Trie trie;
 611     trieobject *trieobj;
 612
 613     if(!PyArg_ParseTuple(args, "O:load", &py_handle))
 614         return NULL;
 615
 616     if(!(trie = Trie_deserialize(_read_from_handle, _read_value_from_handle,
 617                                  py_handle))) {
 618         if(!PyErr_Occurred())
 619             PyErr_SetString(PyExc_RuntimeError,
 620                             "loading failed for some reason");
 621         return NULL;
 622     }
 623
 624     if(!(trieobj = PyObject_New(trieobject, &Trie_Type))) {
 625         Trie_del(trie);
 626         return NULL;
 627     }
 628     trieobj->trie = trie;
 629     return (PyObject *)trieobj;
 630 }
 631
 632 static PyMethodDef trie_methods[] = {
 633     {"trie", trie_trie, METH_VARARGS,
 634      "trie() -> new Trie object."},
 635     {"load", trie_load, METH_VARARGS,
 636      "load(handle) -> trie object"},
 637     {"save", trie_save, METH_VARARGS,
 638      "save(handle, trie), save a trie object to a handle"},
 639     {NULL, NULL, 0, NULL}
 640 };
 641
 642 static char trie__doc__[] =
 643 "\
 644 This module implements a trie data structure.  This allows an O(M)\n\
 645 lookup of a string in a dictionary, where M is the length of the\n\
 646 string.  It also supports approximate matches.\n\
 647 \n\
 648 Functions:\n\
 649 trie    Create a new trie object.\n\
 650 save    Save a trie to a handle.\n\
 651 load    Load a trie from a handle.\n\
 652 \n\
 653 ";
 654
 655 DL_EXPORT(void)
 656 inittrie(void)
 657 {
 658     Trie_Type.ob_type = &PyType_Type;
 659
 660     (void) Py_InitModule3("trie", trie_methods, trie__doc__);
 661 }