pyx/data.py

   1 #!/usr/bin/env python
   2 #
   3 #
   4 # Copyright (C) 2002 Jörg Lehmann <joergl@users.sourceforge.net>
   5 # Copyright (C) 2002 André Wobst <wobsta@users.sourceforge.net>
   6 #
   7 # This file is part of PyX (http://pyx.sourceforge.net/).
   8 #
   9 # PyX is free software; you can redistribute it and/or modify
  10 # it under the terms of the GNU General Public License as published by
  11 # the Free Software Foundation; either version 2 of the License, or
  12 # (at your option) any later version.
  13 #
  14 # PyX is distributed in the hope that it will be useful,
  15 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  16 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  17 # GNU General Public License for more details.
  18 #
  19 # You should have received a copy of the GNU General Public License
  20 # along with PyX; if not, write to the Free Software
  21 # Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
  22
  23
  24 import re, ConfigParser
  25 from mathtree import *
  26
  27
  28 class ColumnError(Exception): pass
  29
  30
  31 ColPattern = re.compile(r"\$-?[0-9]+")
  32
  33 class MathTreeValCol(MathTreeValVar):
  34
  35     def InitByParser(self, arg):
  36         Match = arg.MatchPattern(ColPattern)
  37         if Match:
  38             self.AddArg(Match)
  39             return 1
  40
  41
  42 MathTreeValsWithCol = (MathTreeValConst, MathTreeValVar, MathTreeValCol)
  43
  44
  45 class data:
  46
  47     def __init__(self, titles, data, parser=parser(MathTreeVals=MathTreeValsWithCol), extern=None):
  48         self.titles = titles
  49         self.data = data
  50         self.parser = parser
  51         self.extern = extern
  52
  53     def getcolumnno(self, column):
  54         if self.titles.count(column) == 1:
  55             return self.titles.index(column)
  56         try:
  57             self.titles[column]
  58         except (TypeError, IndexError, ValueError):
  59             raise ColumnError
  60         return column
  61
  62     def getcolumn(self, column):
  63         columnno = self.getcolumnno(column)
  64         return [x[columnno] for x in self.data]
  65
  66     def _addcolumn(self, expression, **columns):
  67         try:
  68             split = expression.rindex("=")
  69         except ValueError:
  70             self.titles.append(None)
  71         else:
  72             self.titles.append(expression[:split])
  73             expression = expression[split+1:]
  74         tree = self.parser.parse(expression)
  75         columnlist = {}
  76         for key in tree.VarList():
  77             if key[0] == "$":
  78                 column = int(key[1:])
  79                 try:
  80                     self.titles[column]
  81                 except:
  82                     raise ColumnError
  83                 columnlist[key] = column
  84             else:
  85                 try:
  86                     columnlist[key] = self.getcolumnno(columns[key])
  87                 except KeyError:
  88                     try:
  89                         columnlist[key] = self.getcolumnno(key)
  90                     except ColumnError, e:
  91                         if self.extern is None or key not in self.extern.keys():
  92                             raise e
  93
  94         varlist = {}
  95         for data in self.data:
  96             try:
  97                 for key in columnlist.keys():
  98                     varlist[key] = float(data[columnlist[key]])
  99             except (TypeError, ValueError):
 100                 data.append(None)
 101             else:
 102                 data.append(tree.Calc(varlist, self.extern))
 103         return columnlist.keys()
 104
 105     def addcolumn(self, expression, **columns):
 106         usedkeys = self._addcolumn(expression, **columns)
 107         unusedkeys = [key for key in columns.keys() if key not in usedkeys]
 108         if len(unusedkeys):
 109             raise KeyError("unused keys %s" % unusedkeys)
 110         return self
 111
 112
 113
 114 class datafile(data):
 115
 116     defaultcommentpattern = re.compile(r"(#+|!+|%+)\s*")
 117     defaultstringpattern = re.compile(r"\"(.*?)\"(\s+|$)")
 118     defaultcolumnpattern = re.compile(r"(.*?)(\s+|$)")
 119
 120     def splitline(self, line, stringpattern, columnpattern, tofloat=1):
 121         result = []
 122         if line.find('"')!=-1 or \
 123            stringpattern is not self.defaultstringpattern or \
 124            columnpattern is not self.defaultcolumnpattern:
 125             while len(line):
 126                 match = stringpattern.match(line)
 127                 if match:
 128                     result.append(match.groups()[0])
 129                     line = line[match.end():]
 130                 else:
 131                     match = columnpattern.match(line)
 132                     if tofloat:
 133                         try:
 134                             result.append(float(match.groups()[0]))
 135                         except (TypeError, ValueError):
 136                             result.append(match.groups()[0])
 137                     else:
 138                         result.append(match.groups()[0])
 139                     line = line[match.end():]
 140         else:
 141             if tofloat:
 142                 try:
 143                     return map(float, line.split())
 144                 except (TypeError, ValueError):
 145                     result = []
 146                     for r in line.split():
 147                         try:
 148                             result.append(float(r))
 149                         except (TypeError, ValueError):
 150                             result.append(r)
 151             else:
 152                 return line.split()
 153
 154         return result
 155
 156     def __init__(self, file, commentpattern=defaultcommentpattern,
 157                              stringpattern=defaultstringpattern,
 158                              columnpattern=defaultcolumnpattern, **args):
 159         try:
 160             file + ''
 161         except TypeError:
 162             pass
 163         else:
 164             file = open(file, "r")
 165         usetitles = []
 166         usedata = []
 167         linenumber = 0
 168         maxcolumns = 0
 169         for line in file.readlines():
 170             line = line.strip()
 171             match = commentpattern.match(line)
 172             if match:
 173                 if not len(usedata):
 174                     newtitles = self.splitline(line[match.end():], stringpattern, columnpattern, tofloat=0)
 175                     if len(newtitles):
 176                         usetitles = newtitles
 177             else:
 178                 linedata = []
 179                 for value in self.splitline(line, stringpattern, columnpattern, tofloat=1):
 180                     linedata.append(value)
 181                 if len(linedata):
 182                     linenumber += 1
 183                     linedata = [linenumber] + linedata
 184                     if len(linedata) > maxcolumns:
 185                         maxcolumns = len(linedata)
 186                     usedata.append(linedata)
 187         if linenumber:
 188             usetitles = [None] + usetitles[:maxcolumns-1]
 189             usetitles += [None] * (maxcolumns - len(usetitles))
 190             for line in usedata:
 191                 line += [None] * (maxcolumns - len(line))
 192         else:
 193             usetitles = []
 194         data.__init__(self, usetitles, usedata, **args)
 195
 196
 197
 198 class sectionfile(data):
 199
 200     def __init__(self, file, sectionstr = "section", **args):
 201         config = ConfigParser.ConfigParser()
 202         config.optionxform = str
 203         try:
 204             file = file + ''
 205         except TypeError:
 206             config.readfp(file)
 207         else:
 208             config.readfp(open(file, "r"))
 209         usedata = []
 210         usetitles = [sectionstr]
 211         sections = config.sections()
 212         sections.sort()
 213         for section in sections:
 214             usedata.append([section] + [None for x in range(len(usetitles) - 1)])
 215             for option in config.options(section):
 216                 if option == sectionstr:
 217                     raise ValueError("'%s' is already used as the section identifier" % sectionstr)
 218                 try:
 219                     index = usetitles.index(option)
 220                 except ValueError:
 221                     index = len(usetitles)
 222                     usetitles.append(option)
 223                     for line in usedata:
 224                         line.append(None)
 225                 value = config.get(section, option)
 226                 try:
 227                     usedata[-1][index] = float(value)
 228                 except (TypeError, ValueError):
 229                     usedata[-1][index] = value
 230         data.__init__(self, usetitles, usedata, **args)
 231