psycopg2/examples/usercast.py

   1 # usercast.py -- example of user defined typecasters
   2 # -*- encoding: latin-1 -*-
   3 #
   4 # Copyright (C) 2001-2005 Federico Di Gregorio  <fog@debian.org>
   5 #
   6 # This program is free software; you can redistribute it and/or modify
   7 # it under the terms of the GNU General Public License as published by the
   8 # Free Software Foundation; either version 2, or (at your option) any later
   9 # version.
  10 #
  11 # This program is distributed in the hope that it will be useful, but
  12 # WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTIBILITY
  13 # or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
  14 # for more details.
  15
  16 ## put in DSN your DSN string
  17
  18 DSN = 'dbname=test'
  19
  20 ## don't modify anything below tis line (except for experimenting)
  21
  22 import sys
  23 import psycopg2
  24 import psycopg2.extensions
  25 import whrandom
  26
  27 # importing psycopg.extras will give us a nice tuple adapter: this is wrong
  28 # because the adapter is meant to be used in SQL IN clauses while we use
  29 # tuples to represent points but it works and the example is about Rect, not
  30 # "Point"
  31 import psycopg2.extras
  32
  33 if len(sys.argv) > 1:
  34     DSN = sys.argv[1]
  35
  36 print "Opening connection using dns:", DSN
  37 conn = psycopg2.connect(DSN)
  38 print "Initial encoding for this connection is", conn.encoding
  39
  40 curs = conn.cursor()
  41 try:
  42     curs.execute("CREATE TABLE test_cast (p1 point, p2 point, b box)")
  43 except:
  44     conn.rollback()
  45     curs.execute("DROP TABLE test_cast")
  46     curs.execute("CREATE TABLE test_cast (p1 point, p2 point, b box)")
  47 conn.commit()
  48
  49 # this is the callable object we use as a typecast (the typecast is
  50 # usually a function, but we use a class, just to demonstrate the
  51 # flexibility of the psycopg casting system
  52
  53 class Rect(object):
  54     """Very simple rectangle.
  55
  56     Note that we use this type as a data holder, as an adapter of itself for
  57     the ISQLQuote protocol used by psycopg's adapt() (see __confrom__ below)
  58     and eventually as a type-caster for the data extracted from the database
  59     (that's why __init__ takes the curs argument.)
  60     """
  61
  62     def __init__(self, s=None, curs=None):
  63         """Init the rectangle from the optional string s."""
  64         self.x = self.y = self.width = self.height = 0.0
  65         if s: self.from_string(s)
  66
  67     def __conform__(self, proto):
  68         """This is a terrible hack, just ignore proto and return self."""
  69         if proto == psycopg2.extensions.ISQLQuote:
  70             return self
  71
  72     def from_points(self, x0, y0, x1, y1):
  73         """Init the rectangle from points."""
  74         if x0 > x1: (x0, x1) = (x1, x0)
  75         if y0 > y1: (y0, y1) = (y1, y0)
  76         self.x = x0
  77         self.y = y0
  78         self.width = x1 - x0
  79         self.height = y1 - y0
  80
  81     def from_string(self, s):
  82         """Init the rectangle from a string."""
  83         seq = eval(s)
  84         self.from_points(seq[0][0], seq[0][1], seq[1][0], seq[1][1])
  85
  86     def getquoted(self):
  87         """Format self as a string usable by the db to represent a box."""
  88         s = "'((%d,%d),(%d,%d))'" % (
  89             self.x, self.y, self.x + self.width, self.y + self.height)
  90         return s
  91
  92     def show(self):
  93         """Format a description of the box."""
  94         s = "X: %d\tY: %d\tWidth: %d\tHeight: %d" % (
  95             self.x, self.y, self.width, self.height)
  96         return s
  97
  98 # here we select from the empty table, just to grab the description
  99 curs.execute("SELECT b FROM test_cast WHERE 0=1")
 100 boxoid = curs.description[0][1]
 101 print "Oid for the box datatype is", boxoid
 102
 103 # and build the user cast object
 104 BOX = psycopg2.extensions.new_type((boxoid,), "BOX", Rect)
 105 psycopg2.extensions.register_type(BOX)
 106
 107 # now insert 100 random data (2 points and a box in each row)
 108 for i in range(100):
 109     p1 = (whrandom.randint(0,100), whrandom.randint(0,100))
 110     p2 = (whrandom.randint(0,100), whrandom.randint(0,100))
 111     b = Rect()
 112     b.from_points(whrandom.randint(0,100), whrandom.randint(0,100),
 113                   whrandom.randint(0,100), whrandom.randint(0,100))
 114     curs.execute("INSERT INTO test_cast VALUES ('%(p1)s', '%(p2)s', %(box)s)",
 115                  {'box':b, 'p1':p1, 'p2':p2})
 116 print "Added 100 boxed to the database"
 117
 118 # select and print all boxes with at least one point inside
 119 curs.execute("SELECT b FROM test_cast WHERE p1 @ b OR p2 @ b")
 120 boxes = curs.fetchall()
 121 print "Found %d boxes with at least a point inside:" % len(boxes)
 122 for box in boxes:
 123     print " ", box[0].show()
 124
 125 curs.execute("DROP TABLE test_cast")
 126 conn.commit()