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 import helper, mathtree
  26
  27
  28 class ColumnError(Exception): pass
  29
  30
  31 ColPattern = re.compile(r"\$-?[0-9]+")
  32
  33 class MathTreeValCol(mathtree.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 = (mathtree.MathTreeValConst,
  43                        mathtree.MathTreeValVar,
  44                        MathTreeValCol)
  45
  46
  47 class _data:
  48
  49     def __init__(self, titles, data, parser=mathtree.parser(MathTreeVals=MathTreeValsWithCol)):
  50         self.titles = titles
  51         self.data = data
  52         self.parser = parser
  53
  54     def getcolumnno(self, column):
  55         if self.titles.count(column) == 1:
  56             return self.titles.index(column)
  57         try:
  58             self.titles[column]
  59         except (TypeError, IndexError, ValueError):
  60             raise ColumnError
  61         return column
  62
  63     def getcolumn(self, column):
  64         columnno = self.getcolumnno(column)
  65         return [x[columnno] for x in self.data]
  66
  67     def addcolumn(self, expression, context={}):
  68         try:
  69             split = expression.rindex("=")
  70         except ValueError:
  71             self.titles.append(None)
  72         else:
  73             self.titles.append(expression[:split])
  74             expression = expression[split+1:]
  75         tree = self.parser.parse(expression)
  76         columnlist = {}
  77         for key in tree.VarList():
  78             if key[0] == "$":
  79                 column = int(key[1:])
  80                 try:
  81                     self.titles[column]
  82                 except:
  83                     raise ColumnError
  84                 columnlist[key] = column
  85             else:
  86                 try:
  87                     columnlist[key] = self.getcolumnno(key)
  88                 except ColumnError, e:
  89                     if key not in context.keys():
  90                         raise e
  91
  92         varlist = context.copy()
  93         for data in self.data:
  94             try:
  95                 for key in columnlist.keys():
  96                     varlist[key] = float(data[columnlist[key]])
  97             except (TypeError, ValueError):
  98                 data.append(None)
  99             else:
 100                 data.append(tree.Calc(**varlist))
 101
 102
 103 class data(_data):
 104
 105     def __init__(self, titles=[], data=[], maxcolumns=helper.nodefault, **kwargs):
 106         if len(data):
 107             if maxcolumns is helper.nodefault:
 108                 maxcolumns = len(data[0])
 109                 for line in data[1:]:
 110                     if len(line) > maxcolumns:
 111                         maxcolumns = len(line)
 112             titles = titles[:maxcolumns]
 113             titles += [None] * (maxcolumns - len(titles))
 114             for line in data:
 115                 line += [None] * (maxcolumns - len(line))
 116         else:
 117             titles = []
 118         _data.__init__(self, titles, data, **kwargs)
 119
 120
 121 class datafile(data):
 122
 123     defaultcommentpattern = re.compile(r"(#+|!+|%+)\s*")
 124     defaultstringpattern = re.compile(r"\"(.*?)\"(\s+|$)")
 125     defaultcolumnpattern = re.compile(r"(.*?)(\s+|$)")
 126
 127     def splitline(self, line, stringpattern, columnpattern, tofloat=1):
 128         result = []
 129         if line.find('"')!=-1 or \
 130            stringpattern is not self.defaultstringpattern or \
 131            columnpattern is not self.defaultcolumnpattern:
 132             while len(line):
 133                 match = stringpattern.match(line)
 134                 if match:
 135                     result.append(match.groups()[0])
 136                     line = line[match.end():]
 137                 else:
 138                     match = columnpattern.match(line)
 139                     if tofloat:
 140                         try:
 141                             result.append(float(match.groups()[0]))
 142                         except (TypeError, ValueError):
 143                             result.append(match.groups()[0])
 144                     else:
 145                         result.append(match.groups()[0])
 146                     line = line[match.end():]
 147         else:
 148             if tofloat:
 149                 try:
 150                     return map(float, line.split())
 151                 except (TypeError, ValueError):
 152                     result = []
 153                     for r in line.split():
 154                         try:
 155                             result.append(float(r))
 156                         except (TypeError, ValueError):
 157                             result.append(r)
 158             else:
 159                 return line.split()
 160
 161         return result
 162
 163     def __init__(self, file, commentpattern=defaultcommentpattern,
 164                              stringpattern=defaultstringpattern,
 165                              columnpattern=defaultcolumnpattern,
 166                              skiphead=0, skiptail=0, every=1,**kwargs):
 167         if helper.isstring(file):
 168             file = open(file, "r")
 169         usetitles = []
 170         usedata = []
 171         linenumber = 0
 172         maxcolumns = 0
 173         for line in file.readlines():
 174             line = line.strip()
 175             match = commentpattern.match(line)
 176             if match:
 177                 if not len(usedata):
 178                     newtitles = self.splitline(line[match.end():], stringpattern, columnpattern, tofloat=0)
 179                     if len(newtitles):
 180                         usetitles = newtitles
 181             else:
 182                 linedata = []
 183                 for value in self.splitline(line, stringpattern, columnpattern, tofloat=1):
 184                     linedata.append(value)
 185                 if len(linedata):
 186                     if linenumber >= skiphead and not ((linenumber - skiphead) % every):
 187                         linedata = [linenumber] + linedata
 188                         if len(linedata) > maxcolumns:
 189                             maxcolumns = len(linedata)
 190                         usedata.append(linedata)
 191                     linenumber += 1
 192         if skiptail:
 193             del usedata[-skiptail:]
 194         data.__init__(self, titles=[None] + usetitles, data=usedata, maxcolumns=maxcolumns, **kwargs)
 195
 196
 197
 198 class sectionfile(_data):
 199
 200     def __init__(self, file, sectionstr = "section", **kwargs):
 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, **kwargs)
 231