├── .github └── FUNDING.yml ├── 1561006651468.png ├── 1561006737134.png ├── 1561006887117.png ├── 17cc62c98820974f8c759dc086dd5acb.png ├── 28069d48cf3f357dd83e42406e10d980.png ├── README.md ├── mipsAudit.png ├── mipsAudit.py ├── prettytable.py └── prettytable.pyc /.github/FUNDING.yml: -------------------------------------------------------------------------------- 1 | # These are supported funding model platforms 2 | 3 | github: # Replace with up to 4 GitHub Sponsors-enabled usernames e.g., [user1, user2] 4 | patreon: # 5 | open_collective: # Replace with a single Open Collective username 6 | ko_fi: # Replace with a single Ko-fi username 7 | tidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel 8 | community_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry 9 | liberapay: # Replace with a single Liberapay username 10 | issuehunt: # Replace with a single IssueHunt username 11 | otechie: # Replace with a single Otechie username 12 | custom: ['http://pic.giantbranch.cn/pic/1551450728861.jpg'] 13 | -------------------------------------------------------------------------------- /1561006651468.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/giantbranch/mipsAudit/a8e5576461a75c92f4ac7f8ecaa9ee096b8db371/1561006651468.png -------------------------------------------------------------------------------- /1561006737134.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/giantbranch/mipsAudit/a8e5576461a75c92f4ac7f8ecaa9ee096b8db371/1561006737134.png -------------------------------------------------------------------------------- /1561006887117.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/giantbranch/mipsAudit/a8e5576461a75c92f4ac7f8ecaa9ee096b8db371/1561006887117.png -------------------------------------------------------------------------------- /17cc62c98820974f8c759dc086dd5acb.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/giantbranch/mipsAudit/a8e5576461a75c92f4ac7f8ecaa9ee096b8db371/17cc62c98820974f8c759dc086dd5acb.png -------------------------------------------------------------------------------- /28069d48cf3f357dd83e42406e10d980.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/giantbranch/mipsAudit/a8e5576461a75c92f4ac7f8ecaa9ee096b8db371/28069d48cf3f357dd83e42406e10d980.png -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # IDAPython mipsAudit 2 | 3 | ## 简介 4 | 5 | 这是一个简单的IDAPython脚本。 6 | 7 | 进一步来说是MIPS静态汇编审计辅助脚本。 8 | 9 | 可能会有bug，欢迎大家完善。 10 | 11 | 12 | 13 | ## 功能 14 | 15 | 辅助脚本功能如下： 16 | 17 | 1. 找到危险函数的调用处，并且高亮该行（也可以下断点,这个需要自己去源码看吧） 18 | 19 | 2. 给参数赋值处加上注释 20 | 21 | 3. 最后以表格的形式输出函数名，调用地址，参数，还有当前函数的缓冲区大小 22 | 23 | **大家双击addr那一列的地址，即可跳到对应的地址处** 24 | 25 | ![17cc62c98820974f8c759dc086dd5acb](17cc62c98820974f8c759dc086dd5acb.png) 26 | 27 | ![28069d48cf3f357dd83e42406e10d980](28069d48cf3f357dd83e42406e10d980.png) 28 | 29 | ## 审计的危险函数如下 30 | 31 | ``` 32 | dangerous_functions = [ 33 | "strcpy", 34 | "strcat", 35 | "sprintf", 36 | "read", 37 | "getenv" 38 | ] 39 | 40 | attention_function = [ 41 | "memcpy", 42 | "strncpy", 43 | "sscanf", 44 | "strncat", 45 | "snprintf", 46 | "vprintf", 47 | "printf" 48 | ] 49 | 50 | command_execution_function = [ 51 | "system", 52 | "execve", 53 | "popen", 54 | "unlink" 55 | ] 56 | ``` 57 | 58 | ## 使用 59 | 60 | File - Script file 61 | 62 | ![1561006651468](./1561006651468.png) 63 | 64 | 选择mipsAudit.py 65 | 66 | ![1561006737134](./1561006737134.png) 67 | 68 | 即可看到效果 69 | 70 | ![mipsAudit](./mipsAudit.png) 71 | 72 | 双击地址即可跳到对应的代码处 73 | 74 | ![1561006887117](./1561006887117.png) 75 | -------------------------------------------------------------------------------- /mipsAudit.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/giantbranch/mipsAudit/a8e5576461a75c92f4ac7f8ecaa9ee096b8db371/mipsAudit.png -------------------------------------------------------------------------------- /mipsAudit.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # reference 4 | # 《ida pro 权威指南》 5 | # 《python 灰帽子》 6 | # 《家用路由器0day漏洞挖掘》 7 | # https://github.com/wangzery/SearchOverflow/blob/master/SearchOverflow.py 8 | 9 | from idaapi import * 10 | from prettytable import PrettyTable 11 | 12 | 13 | DEBUG = True 14 | 15 | # fgetc,fgets,fread,fprintf, 16 | # vspritnf 17 | 18 | # set function_name 19 | dangerous_functions = [ 20 | "strcpy", 21 | "strcat", 22 | "sprintf", 23 | "read", 24 | "getenv" 25 | ] 26 | 27 | attention_function = [ 28 | "memcpy", 29 | "strncpy", 30 | "sscanf", 31 | "strncat", 32 | "snprintf", 33 | "vprintf", 34 | "printf" 35 | ] 36 | 37 | command_execution_function = [ 38 | "system", 39 | "execve", 40 | "popen", 41 | "unlink" 42 | ] 43 | 44 | # describe arg num of function 45 | 46 | one_arg_function = [ 47 | "getenv", 48 | "system", 49 | "unlink" 50 | ] 51 | 52 | two_arg_function = [ 53 | "strcpy", 54 | "strcat", 55 | "popen" 56 | ] 57 | 58 | three_arg_function = [ 59 | "strncpy", 60 | "strncat", 61 | "memcpy", 62 | "execve", 63 | "read" 64 | ] 65 | 66 | format_function_offset_dict = { 67 | "sprintf":1, 68 | "sscanf":1, 69 | "snprintf":2, 70 | "vprintf":0, 71 | "printf":0 72 | } 73 | 74 | def printFunc(func_name): 75 | string1 = "========================================" 76 | string2 = "========== Aduiting " + func_name + " " 77 | strlen = len(string1) - len(string2) 78 | return string1 + "\n" + string2 + '=' * strlen + "\n" + string1 79 | 80 | def getFuncAddr(func_name): 81 | func_addr = LocByName(func_name) 82 | if func_addr != BADADDR: 83 | print printFunc(func_name) 84 | # print func_name + " Addr : 0x %x" % func_addr 85 | return func_addr 86 | return False 87 | 88 | def getFormatString(addr): 89 | op_num = 1 90 | # GetOpType Return value 91 | #define o_void 0 // No Operand ---------- 92 | #define o_reg 1 // General Register (al, ax, es, ds...) reg 93 | #define o_mem 2 // Direct Memory Reference (DATA) addr 94 | #define o_phrase 3 // Memory Ref [Base Reg + Index Reg] phrase 95 | #define o_displ 4 // Memory Reg [Base Reg + Index Reg + Displacement] phrase+addr 96 | #define o_imm 5 // Immediate Value value 97 | #define o_far 6 // Immediate Far Address (CODE) addr 98 | #define o_near 7 // Immediate Near Address (CODE) addr 99 | #define o_idpspec0 8 // IDP specific type 100 | #define o_idpspec1 9 // IDP specific type 101 | #define o_idpspec2 10 // IDP specific type 102 | #define o_idpspec3 11 // IDP specific type 103 | #define o_idpspec4 12 // IDP specific type 104 | #define o_idpspec5 13 // IDP specific type 105 | # 如果第二个不是立即数则下一个 106 | if(GetOpType(addr ,op_num) != 5): 107 | op_num = op_num + 1 108 | if GetOpType(addr ,op_num) != 5: 109 | return "get fail" 110 | op_string = GetOpnd(addr, op_num).split(" ")[0].split("+")[0].split("-")[0].replace("(", "") 111 | string_addr = LocByName(op_string) 112 | if string_addr == BADADDR: 113 | return "get fail" 114 | string = str(GetString(string_addr)) 115 | return [string_addr, string] 116 | 117 | 118 | def getArgAddr(start_addr, regNum): 119 | mipscondition = ["bn", "be" , "bg", "bl"] 120 | scan_deep = 50 121 | count = 0 122 | reg = "$a" + str(regNum) 123 | # try to get in the next 124 | next_addr = Rfirst(start_addr) 125 | if next_addr != BADADDR and reg == GetOpnd(next_addr, 0): 126 | return next_addr 127 | # try to get before 128 | before_addr = RfirstB(start_addr) 129 | while before_addr != BADADDR: 130 | if reg == GetOpnd(before_addr, 0): 131 | Mnemonics = GetMnem(before_addr) 132 | if Mnemonics[0:2] in mipscondition: 133 | pass 134 | elif Mnemonics[0:1] == "j": 135 | pass 136 | else: 137 | return before_addr 138 | count = count + 1 139 | if count > scan_deep: 140 | break 141 | before_addr = RfirstB(before_addr) 142 | return BADADDR 143 | 144 | 145 | def getArg(start_addr, regNum): 146 | mipsmov = ["move", "lw", "li", "lb", "lui", "lhu", "lbu", "la"] 147 | arg_addr = getArgAddr(start_addr, regNum) 148 | if arg_addr != BADADDR: 149 | Mnemonics = GetMnem(arg_addr) 150 | if Mnemonics[0:3] == "add": 151 | if GetOpnd(arg_addr, 2) == "": 152 | arg = GetOpnd(arg_addr, 0) + "+" + GetOpnd(arg_addr, 1) 153 | else: 154 | arg = GetOpnd(arg_addr, 1) + "+" + GetOpnd(arg_addr, 2) 155 | elif Mnemonics[0:3] == "sub": 156 | if GetOpnd(arg_addr, 2) == "": 157 | arg = GetOpnd(arg_addr, 0) + "-" + GetOpnd(arg_addr, 1) 158 | else: 159 | arg = GetOpnd(arg_addr, 1) + "-" + GetOpnd(arg_addr, 2) 160 | elif Mnemonics in mipsmov: 161 | arg = GetOpnd(arg_addr, 1) 162 | else: 163 | arg = GetDisasm(arg_addr).split("#")[0] 164 | MakeComm(arg_addr, "addr: 0x%x " % start_addr + "-------> arg" + str((int(regNum)+1)) + " : " + arg) 165 | return arg 166 | else: 167 | return "get fail" 168 | 169 | def audit(func_name): 170 | func_addr = getFuncAddr(func_name) 171 | if func_addr == False: 172 | return False 173 | 174 | # get arg num and set table 175 | if func_name in one_arg_function: 176 | arg_num = 1 177 | elif func_name in two_arg_function: 178 | arg_num = 2 179 | elif func_name in three_arg_function: 180 | arg_num = 3 181 | elif func_name in format_function_offset_dict: 182 | arg_num = format_function_offset_dict[func_name] + 1 183 | else: 184 | print "The %s function didn't write in the describe arg num of function array,please add it to,such as add to `two_arg_function` arary" % func_name 185 | return 186 | # mispcall = ["jal", "jalr", "bal", "jr"] 187 | table_head = ["func_name", "addr"] 188 | for num in xrange(0,arg_num): 189 | table_head.append("arg"+str(num+1)) 190 | if func_name in format_function_offset_dict: 191 | table_head.append("format&value[string_addr, num of '%', fmt_arg...]") 192 | table_head.append("local_buf_size") 193 | table = PrettyTable(table_head) 194 | 195 | # get first call 196 | call_addr = RfirstB(func_addr) 197 | while call_addr != BADADDR: 198 | # set color ———— green (red=0x0000ff,blue = 0xff0000) 199 | SetColor(call_addr, CIC_ITEM, 0x00ff00) 200 | # set break point 201 | # AddBpt(call_addr) 202 | # DelBpt(call_addr) 203 | 204 | # if you want to use condition 205 | # SetBptCnd(ea, 'strstr(GetString(Dword(esp+4),-1, 0), "SAEXT.DLL") != -1') 206 | Mnemonics = GetMnem(call_addr) 207 | # print "Mnemonics : %s" % Mnemonics 208 | # if Mnemonics in mispcall: 209 | if Mnemonics[0:1] == "j" or Mnemonics[0:1] == "b": 210 | # print func + " addr : 0x%x" % call_addr 211 | if func_name in format_function_offset_dict: 212 | info = auditFormat(call_addr, func_name, arg_num) 213 | else: 214 | info = auditAddr(call_addr, func_name, arg_num) 215 | table.add_row(info) 216 | call_addr = RnextB(func_addr, call_addr) 217 | print table 218 | # data_addr = DfirstB(func_addr) 219 | # while data_addr != BADADDR: 220 | # Mnemonics = GetMnem(data_addr) 221 | # if DEBUG: 222 | # print "Data Mnemonics : %s" % GetMnem(data_addr) 223 | # print "Data addr : 0x %s" % data_addr 224 | # data_addr = DnextB(func_addr, data_addr) 225 | 226 | def auditAddr(call_addr, func_name, arg_num): 227 | addr = "0x%x" % call_addr 228 | ret_list = [func_name, addr] 229 | # local buf size 230 | local_buf_size = GetFunctionAttr(call_addr , FUNCATTR_FRSIZE) 231 | if local_buf_size == BADADDR : 232 | local_buf_size = "get fail" 233 | else: 234 | local_buf_size = "0x%x" % local_buf_size 235 | # get arg 236 | for num in xrange(0,arg_num): 237 | ret_list.append(getArg(call_addr, num)) 238 | ret_list.append(local_buf_size) 239 | return ret_list 240 | 241 | def auditFormat(call_addr, func_name, arg_num): 242 | addr = "0x%x" % call_addr 243 | ret_list = [func_name, addr] 244 | # local buf size 245 | local_buf_size = GetFunctionAttr(call_addr , FUNCATTR_FRSIZE) 246 | if local_buf_size == BADADDR : 247 | local_buf_size = "get fail" 248 | else: 249 | local_buf_size = "0x%x" % local_buf_size 250 | # get arg 251 | for num in xrange(0,arg_num): 252 | ret_list.append(getArg(call_addr, num)) 253 | arg_addr = getArgAddr(call_addr, format_function_offset_dict[func_name]) 254 | string_and_addr = getFormatString(arg_addr) 255 | format_and_value = [] 256 | if string_and_addr == "get fail": 257 | ret_list.append("get fail") 258 | else: 259 | string_addr = "0x%x" % string_and_addr[0] 260 | format_and_value.append(string_addr) 261 | string = string_and_addr[1] 262 | fmt_num = string.count("%") 263 | format_and_value.append(fmt_num) 264 | # mips arg reg is from a0 to a3 265 | if fmt_num > 3: 266 | fmt_num = fmt_num - format_function_offset_dict[func_name] - 1 267 | for num in xrange(0,fmt_num): 268 | if arg_num + num > 3: 269 | break 270 | format_and_value.append(getArg(call_addr, arg_num + num)) 271 | ret_list.append(format_and_value) 272 | # format_string = str(getFormatString(arg_addr)[1]) 273 | 274 | # print " format String: " + format_string 275 | # ret_list.append([string_addr]) 276 | ret_list.append(local_buf_size) 277 | return ret_list 278 | 279 | def mipsAudit(): 280 | # the word create with figlet 281 | start = ''' 282 | _ _ _ _ _ 283 | _ __ ___ (_)_ __ ___ / \ _ _ __| (_) |_ 284 | | '_ ` _ \| | '_ \/ __| / _ \| | | |/ _` | | __| 285 | | | | | | | | |_) \__ \/ ___ \ |_| | (_| | | |_ 286 | |_| |_| |_|_| .__/|___/_/ \_\__,_|\__,_|_|\__| 287 | |_| 288 | code by giantbranch 2018.05 289 | ''' 290 | print start 291 | print "Auditing dangerous functions ......" 292 | for func_name in dangerous_functions: 293 | audit(func_name) 294 | 295 | print "Auditing attention function ......" 296 | for func_name in attention_function: 297 | audit(func_name) 298 | 299 | print "Auditing command execution function ......" 300 | for func_name in command_execution_function: 301 | audit(func_name) 302 | 303 | print "Finished! Enjoy the result ~" 304 | 305 | # 判断架构的代码，以后或许用得上 306 | # info = idaapi.get_inf_structure() 307 | 308 | # if info.is_64bit(): 309 | # bits = 64 310 | # elif info.is_32bit(): 311 | # bits = 32 312 | # else: 313 | # bits = 16 314 | 315 | # try: 316 | # is_be = info.is_be() 317 | # except: 318 | # is_be = info.mf 319 | # endian = "big" if is_be else "little" 320 | 321 | # print 'Processor: {}, {}bit, {} endian'.format(info.procName, bits, endian) 322 | # # Result: Processor: mipsr, 32bit, big endian 323 | 324 | mipsAudit() -------------------------------------------------------------------------------- /prettytable.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # 3 | # Copyright (c) 2009-2013, Luke Maurits 4 | # All rights reserved. 5 | # With contributions from: 6 | # * Chris Clark 7 | # * Klein Stephane 8 | # 9 | # Redistribution and use in source and binary forms, with or without 10 | # modification, are permitted provided that the following conditions are met: 11 | # 12 | # * Redistributions of source code must retain the above copyright notice, 13 | # this list of conditions and the following disclaimer. 14 | # * Redistributions in binary form must reproduce the above copyright notice, 15 | # this list of conditions and the following disclaimer in the documentation 16 | # and/or other materials provided with the distribution. 17 | # * The name of the author may not be used to endorse or promote products 18 | # derived from this software without specific prior written permission. 19 | # 20 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 21 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 22 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 23 | # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE 24 | # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 25 | # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 26 | # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 27 | # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 28 | # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 29 | # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 30 | # POSSIBILITY OF SUCH DAMAGE. 31 | 32 | __version__ = "0.7.2" 33 | 34 | import copy 35 | import csv 36 | import random 37 | import re 38 | import sys 39 | import textwrap 40 | import itertools 41 | import unicodedata 42 | 43 | py3k = sys.version_info[0] >= 3 44 | if py3k: 45 | unicode = str 46 | basestring = str 47 | itermap = map 48 | iterzip = zip 49 | uni_chr = chr 50 | from html.parser import HTMLParser 51 | else: 52 | itermap = itertools.imap 53 | iterzip = itertools.izip 54 | uni_chr = unichr 55 | from HTMLParser import HTMLParser 56 | 57 | if py3k and sys.version_info[1] >= 2: 58 | from html import escape 59 | else: 60 | from cgi import escape 61 | 62 | # hrule styles 63 | FRAME = 0 64 | ALL = 1 65 | NONE = 2 66 | HEADER = 3 67 | 68 | # Table styles 69 | DEFAULT = 10 70 | MSWORD_FRIENDLY = 11 71 | PLAIN_COLUMNS = 12 72 | RANDOM = 20 73 | 74 | _re = re.compile("\033\[[0-9;]*m") 75 | 76 | def _get_size(text): 77 | lines = text.split("\n") 78 | height = len(lines) 79 | width = max([_str_block_width(line) for line in lines]) 80 | return (width, height) 81 | 82 | class PrettyTable(object): 83 | 84 | def __init__(self, field_names=None, **kwargs): 85 | 86 | """Return a new PrettyTable instance 87 | 88 | Arguments: 89 | 90 | encoding - Unicode encoding scheme used to decode any encoded input 91 | field_names - list or tuple of field names 92 | fields - list or tuple of field names to include in displays 93 | start - index of first data row to include in output 94 | end - index of last data row to include in output PLUS ONE (list slice style) 95 | header - print a header showing field names (True or False) 96 | header_style - stylisation to apply to field names in header ("cap", "title", "upper", "lower" or None) 97 | border - print a border around the table (True or False) 98 | hrules - controls printing of horizontal rules after rows. Allowed values: FRAME, HEADER, ALL, NONE 99 | vrules - controls printing of vertical rules between columns. Allowed values: FRAME, ALL, NONE 100 | int_format - controls formatting of integer data 101 | float_format - controls formatting of floating point data 102 | padding_width - number of spaces on either side of column data (only used if left and right paddings are None) 103 | left_padding_width - number of spaces on left hand side of column data 104 | right_padding_width - number of spaces on right hand side of column data 105 | vertical_char - single character string used to draw vertical lines 106 | horizontal_char - single character string used to draw horizontal lines 107 | junction_char - single character string used to draw line junctions 108 | sortby - name of field to sort rows by 109 | sort_key - sorting key function, applied to data points before sorting 110 | valign - default valign for each row (None, "t", "m" or "b") 111 | reversesort - True or False to sort in descending or ascending order""" 112 | 113 | self.encoding = kwargs.get("encoding", "UTF-8") 114 | 115 | # Data 116 | self._field_names = [] 117 | self._align = {} 118 | self._valign = {} 119 | self._max_width = {} 120 | self._rows = [] 121 | if field_names: 122 | self.field_names = field_names 123 | else: 124 | self._widths = [] 125 | 126 | # Options 127 | self._options = "start end fields header border sortby reversesort sort_key attributes format hrules vrules".split() 128 | self._options.extend("int_format float_format padding_width left_padding_width right_padding_width".split()) 129 | self._options.extend("vertical_char horizontal_char junction_char header_style valign xhtml print_empty".split()) 130 | for option in self._options: 131 | if option in kwargs: 132 | self._validate_option(option, kwargs[option]) 133 | else: 134 | kwargs[option] = None 135 | 136 | self._start = kwargs["start"] or 0 137 | self._end = kwargs["end"] or None 138 | self._fields = kwargs["fields"] or None 139 | 140 | if kwargs["header"] in (True, False): 141 | self._header = kwargs["header"] 142 | else: 143 | self._header = True 144 | self._header_style = kwargs["header_style"] or None 145 | if kwargs["border"] in (True, False): 146 | self._border = kwargs["border"] 147 | else: 148 | self._border = True 149 | self._hrules = kwargs["hrules"] or FRAME 150 | self._vrules = kwargs["vrules"] or ALL 151 | 152 | self._sortby = kwargs["sortby"] or None 153 | if kwargs["reversesort"] in (True, False): 154 | self._reversesort = kwargs["reversesort"] 155 | else: 156 | self._reversesort = False 157 | self._sort_key = kwargs["sort_key"] or (lambda x: x) 158 | 159 | self._int_format = kwargs["int_format"] or {} 160 | self._float_format = kwargs["float_format"] or {} 161 | self._padding_width = kwargs["padding_width"] or 1 162 | self._left_padding_width = kwargs["left_padding_width"] or None 163 | self._right_padding_width = kwargs["right_padding_width"] or None 164 | 165 | self._vertical_char = kwargs["vertical_char"] or self._unicode("|") 166 | self._horizontal_char = kwargs["horizontal_char"] or self._unicode("-") 167 | self._junction_char = kwargs["junction_char"] or self._unicode("+") 168 | 169 | if kwargs["print_empty"] in (True, False): 170 | self._print_empty = kwargs["print_empty"] 171 | else: 172 | self._print_empty = True 173 | self._format = kwargs["format"] or False 174 | self._xhtml = kwargs["xhtml"] or False 175 | self._attributes = kwargs["attributes"] or {} 176 | 177 | def _unicode(self, value): 178 | if not isinstance(value, basestring): 179 | value = str(value) 180 | if not isinstance(value, unicode): 181 | value = unicode(value, self.encoding, "strict") 182 | return value 183 | 184 | def _justify(self, text, width, align): 185 | excess = width - _str_block_width(text) 186 | if align == "l": 187 | return text + excess * " " 188 | elif align == "r": 189 | return excess * " " + text 190 | else: 191 | if excess % 2: 192 | # Uneven padding 193 | # Put more space on right if text is of odd length... 194 | if _str_block_width(text) % 2: 195 | return (excess//2)*" " + text + (excess//2 + 1)*" " 196 | # and more space on left if text is of even length 197 | else: 198 | return (excess//2 + 1)*" " + text + (excess//2)*" " 199 | # Why distribute extra space this way? To match the behaviour of 200 | # the inbuilt str.center() method. 201 | else: 202 | # Equal padding on either side 203 | return (excess//2)*" " + text + (excess//2)*" " 204 | 205 | def __getattr__(self, name): 206 | 207 | if name == "rowcount": 208 | return len(self._rows) 209 | elif name == "colcount": 210 | if self._field_names: 211 | return len(self._field_names) 212 | elif self._rows: 213 | return len(self._rows[0]) 214 | else: 215 | return 0 216 | else: 217 | raise AttributeError(name) 218 | 219 | def __getitem__(self, index): 220 | 221 | new = PrettyTable() 222 | new.field_names = self.field_names 223 | for attr in self._options: 224 | setattr(new, "_"+attr, getattr(self, "_"+attr)) 225 | setattr(new, "_align", getattr(self, "_align")) 226 | if isinstance(index, slice): 227 | for row in self._rows[index]: 228 | new.add_row(row) 229 | elif isinstance(index, int): 230 | new.add_row(self._rows[index]) 231 | else: 232 | raise Exception("Index %s is invalid, must be an integer or slice" % str(index)) 233 | return new 234 | 235 | if py3k: 236 | def __str__(self): 237 | return self.__unicode__() 238 | else: 239 | def __str__(self): 240 | return self.__unicode__().encode(self.encoding) 241 | 242 | def __unicode__(self): 243 | return self.get_string() 244 | 245 | ############################## 246 | # ATTRIBUTE VALIDATORS # 247 | ############################## 248 | 249 | # The method _validate_option is all that should be used elsewhere in the code base to validate options. 250 | # It will call the appropriate validation method for that option. The individual validation methods should 251 | # never need to be called directly (although nothing bad will happen if they *are*). 252 | # Validation happens in TWO places. 253 | # Firstly, in the property setters defined in the ATTRIBUTE MANAGMENT section. 254 | # Secondly, in the _get_options method, where keyword arguments are mixed with persistent settings 255 | 256 | def _validate_option(self, option, val): 257 | if option in ("field_names"): 258 | self._validate_field_names(val) 259 | elif option in ("start", "end", "max_width", "padding_width", "left_padding_width", "right_padding_width", "format"): 260 | self._validate_nonnegative_int(option, val) 261 | elif option in ("sortby"): 262 | self._validate_field_name(option, val) 263 | elif option in ("sort_key"): 264 | self._validate_function(option, val) 265 | elif option in ("hrules"): 266 | self._validate_hrules(option, val) 267 | elif option in ("vrules"): 268 | self._validate_vrules(option, val) 269 | elif option in ("fields"): 270 | self._validate_all_field_names(option, val) 271 | elif option in ("header", "border", "reversesort", "xhtml", "print_empty"): 272 | self._validate_true_or_false(option, val) 273 | elif option in ("header_style"): 274 | self._validate_header_style(val) 275 | elif option in ("int_format"): 276 | self._validate_int_format(option, val) 277 | elif option in ("float_format"): 278 | self._validate_float_format(option, val) 279 | elif option in ("vertical_char", "horizontal_char", "junction_char"): 280 | self._validate_single_char(option, val) 281 | elif option in ("attributes"): 282 | self._validate_attributes(option, val) 283 | else: 284 | raise Exception("Unrecognised option: %s!" % option) 285 | 286 | def _validate_field_names(self, val): 287 | # Check for appropriate length 288 | if self._field_names: 289 | try: 290 | assert len(val) == len(self._field_names) 291 | except AssertionError: 292 | raise Exception("Field name list has incorrect number of values, (actual) %d!=%d (expected)" % (len(val), len(self._field_names))) 293 | if self._rows: 294 | try: 295 | assert len(val) == len(self._rows[0]) 296 | except AssertionError: 297 | raise Exception("Field name list has incorrect number of values, (actual) %d!=%d (expected)" % (len(val), len(self._rows[0]))) 298 | # Check for uniqueness 299 | try: 300 | assert len(val) == len(set(val)) 301 | except AssertionError: 302 | raise Exception("Field names must be unique!") 303 | 304 | def _validate_header_style(self, val): 305 | try: 306 | assert val in ("cap", "title", "upper", "lower", None) 307 | except AssertionError: 308 | raise Exception("Invalid header style, use cap, title, upper, lower or None!") 309 | 310 | def _validate_align(self, val): 311 | try: 312 | assert val in ["l","c","r"] 313 | except AssertionError: 314 | raise Exception("Alignment %s is invalid, use l, c or r!" % val) 315 | 316 | def _validate_valign(self, val): 317 | try: 318 | assert val in ["t","m","b",None] 319 | except AssertionError: 320 | raise Exception("Alignment %s is invalid, use t, m, b or None!" % val) 321 | 322 | def _validate_nonnegative_int(self, name, val): 323 | try: 324 | assert int(val) >= 0 325 | except AssertionError: 326 | raise Exception("Invalid value for %s: %s!" % (name, self._unicode(val))) 327 | 328 | def _validate_true_or_false(self, name, val): 329 | try: 330 | assert val in (True, False) 331 | except AssertionError: 332 | raise Exception("Invalid value for %s! Must be True or False." % name) 333 | 334 | def _validate_int_format(self, name, val): 335 | if val == "": 336 | return 337 | try: 338 | assert type(val) in (str, unicode) 339 | assert val.isdigit() 340 | except AssertionError: 341 | raise Exception("Invalid value for %s! Must be an integer format string." % name) 342 | 343 | def _validate_float_format(self, name, val): 344 | if val == "": 345 | return 346 | try: 347 | assert type(val) in (str, unicode) 348 | assert "." in val 349 | bits = val.split(".") 350 | assert len(bits) <= 2 351 | assert bits[0] == "" or bits[0].isdigit() 352 | assert bits[1] == "" or bits[1].isdigit() 353 | except AssertionError: 354 | raise Exception("Invalid value for %s! Must be a float format string." % name) 355 | 356 | def _validate_function(self, name, val): 357 | try: 358 | assert hasattr(val, "__call__") 359 | except AssertionError: 360 | raise Exception("Invalid value for %s! Must be a function." % name) 361 | 362 | def _validate_hrules(self, name, val): 363 | try: 364 | assert val in (ALL, FRAME, HEADER, NONE) 365 | except AssertionError: 366 | raise Exception("Invalid value for %s! Must be ALL, FRAME, HEADER or NONE." % name) 367 | 368 | def _validate_vrules(self, name, val): 369 | try: 370 | assert val in (ALL, FRAME, NONE) 371 | except AssertionError: 372 | raise Exception("Invalid value for %s! Must be ALL, FRAME, or NONE." % name) 373 | 374 | def _validate_field_name(self, name, val): 375 | try: 376 | assert (val in self._field_names) or (val is None) 377 | except AssertionError: 378 | raise Exception("Invalid field name: %s!" % val) 379 | 380 | def _validate_all_field_names(self, name, val): 381 | try: 382 | for x in val: 383 | self._validate_field_name(name, x) 384 | except AssertionError: 385 | raise Exception("fields must be a sequence of field names!") 386 | 387 | def _validate_single_char(self, name, val): 388 | try: 389 | assert _str_block_width(val) == 1 390 | except AssertionError: 391 | raise Exception("Invalid value for %s! Must be a string of length 1." % name) 392 | 393 | def _validate_attributes(self, name, val): 394 | try: 395 | assert isinstance(val, dict) 396 | except AssertionError: 397 | raise Exception("attributes must be a dictionary of name/value pairs!") 398 | 399 | ############################## 400 | # ATTRIBUTE MANAGEMENT # 401 | ############################## 402 | 403 | def _get_field_names(self): 404 | return self._field_names 405 | """The names of the fields 406 | 407 | Arguments: 408 | 409 | fields - list or tuple of field names""" 410 | def _set_field_names(self, val): 411 | val = [self._unicode(x) for x in val] 412 | self._validate_option("field_names", val) 413 | if self._field_names: 414 | old_names = self._field_names[:] 415 | self._field_names = val 416 | if self._align and old_names: 417 | for old_name, new_name in zip(old_names, val): 418 | self._align[new_name] = self._align[old_name] 419 | for old_name in old_names: 420 | if old_name not in self._align: 421 | self._align.pop(old_name) 422 | else: 423 | for field in self._field_names: 424 | self._align[field] = "c" 425 | if self._valign and old_names: 426 | for old_name, new_name in zip(old_names, val): 427 | self._valign[new_name] = self._valign[old_name] 428 | for old_name in old_names: 429 | if old_name not in self._valign: 430 | self._valign.pop(old_name) 431 | else: 432 | for field in self._field_names: 433 | self._valign[field] = "t" 434 | field_names = property(_get_field_names, _set_field_names) 435 | 436 | def _get_align(self): 437 | return self._align 438 | def _set_align(self, val): 439 | self._validate_align(val) 440 | for field in self._field_names: 441 | self._align[field] = val 442 | align = property(_get_align, _set_align) 443 | 444 | def _get_valign(self): 445 | return self._valign 446 | def _set_valign(self, val): 447 | self._validate_valign(val) 448 | for field in self._field_names: 449 | self._valign[field] = val 450 | valign = property(_get_valign, _set_valign) 451 | 452 | def _get_max_width(self): 453 | return self._max_width 454 | def _set_max_width(self, val): 455 | self._validate_option("max_width", val) 456 | for field in self._field_names: 457 | self._max_width[field] = val 458 | max_width = property(_get_max_width, _set_max_width) 459 | 460 | def _get_fields(self): 461 | """List or tuple of field names to include in displays 462 | 463 | Arguments: 464 | 465 | fields - list or tuple of field names to include in displays""" 466 | return self._fields 467 | def _set_fields(self, val): 468 | self._validate_option("fields", val) 469 | self._fields = val 470 | fields = property(_get_fields, _set_fields) 471 | 472 | def _get_start(self): 473 | """Start index of the range of rows to print 474 | 475 | Arguments: 476 | 477 | start - index of first data row to include in output""" 478 | return self._start 479 | 480 | def _set_start(self, val): 481 | self._validate_option("start", val) 482 | self._start = val 483 | start = property(_get_start, _set_start) 484 | 485 | def _get_end(self): 486 | """End index of the range of rows to print 487 | 488 | Arguments: 489 | 490 | end - index of last data row to include in output PLUS ONE (list slice style)""" 491 | return self._end 492 | def _set_end(self, val): 493 | self._validate_option("end", val) 494 | self._end = val 495 | end = property(_get_end, _set_end) 496 | 497 | def _get_sortby(self): 498 | """Name of field by which to sort rows 499 | 500 | Arguments: 501 | 502 | sortby - field name to sort by""" 503 | return self._sortby 504 | def _set_sortby(self, val): 505 | self._validate_option("sortby", val) 506 | self._sortby = val 507 | sortby = property(_get_sortby, _set_sortby) 508 | 509 | def _get_reversesort(self): 510 | """Controls direction of sorting (ascending vs descending) 511 | 512 | Arguments: 513 | 514 | reveresort - set to True to sort by descending order, or False to sort by ascending order""" 515 | return self._reversesort 516 | def _set_reversesort(self, val): 517 | self._validate_option("reversesort", val) 518 | self._reversesort = val 519 | reversesort = property(_get_reversesort, _set_reversesort) 520 | 521 | def _get_sort_key(self): 522 | """Sorting key function, applied to data points before sorting 523 | 524 | Arguments: 525 | 526 | sort_key - a function which takes one argument and returns something to be sorted""" 527 | return self._sort_key 528 | def _set_sort_key(self, val): 529 | self._validate_option("sort_key", val) 530 | self._sort_key = val 531 | sort_key = property(_get_sort_key, _set_sort_key) 532 | 533 | def _get_header(self): 534 | """Controls printing of table header with field names 535 | 536 | Arguments: 537 | 538 | header - print a header showing field names (True or False)""" 539 | return self._header 540 | def _set_header(self, val): 541 | self._validate_option("header", val) 542 | self._header = val 543 | header = property(_get_header, _set_header) 544 | 545 | def _get_header_style(self): 546 | """Controls stylisation applied to field names in header 547 | 548 | Arguments: 549 | 550 | header_style - stylisation to apply to field names in header ("cap", "title", "upper", "lower" or None)""" 551 | return self._header_style 552 | def _set_header_style(self, val): 553 | self._validate_header_style(val) 554 | self._header_style = val 555 | header_style = property(_get_header_style, _set_header_style) 556 | 557 | def _get_border(self): 558 | """Controls printing of border around table 559 | 560 | Arguments: 561 | 562 | border - print a border around the table (True or False)""" 563 | return self._border 564 | def _set_border(self, val): 565 | self._validate_option("border", val) 566 | self._border = val 567 | border = property(_get_border, _set_border) 568 | 569 | def _get_hrules(self): 570 | """Controls printing of horizontal rules after rows 571 | 572 | Arguments: 573 | 574 | hrules - horizontal rules style. Allowed values: FRAME, ALL, HEADER, NONE""" 575 | return self._hrules 576 | def _set_hrules(self, val): 577 | self._validate_option("hrules", val) 578 | self._hrules = val 579 | hrules = property(_get_hrules, _set_hrules) 580 | 581 | def _get_vrules(self): 582 | """Controls printing of vertical rules between columns 583 | 584 | Arguments: 585 | 586 | vrules - vertical rules style. Allowed values: FRAME, ALL, NONE""" 587 | return self._vrules 588 | def _set_vrules(self, val): 589 | self._validate_option("vrules", val) 590 | self._vrules = val 591 | vrules = property(_get_vrules, _set_vrules) 592 | 593 | def _get_int_format(self): 594 | """Controls formatting of integer data 595 | Arguments: 596 | 597 | int_format - integer format string""" 598 | return self._int_format 599 | def _set_int_format(self, val): 600 | # self._validate_option("int_format", val) 601 | for field in self._field_names: 602 | self._int_format[field] = val 603 | int_format = property(_get_int_format, _set_int_format) 604 | 605 | def _get_float_format(self): 606 | """Controls formatting of floating point data 607 | Arguments: 608 | 609 | float_format - floating point format string""" 610 | return self._float_format 611 | def _set_float_format(self, val): 612 | # self._validate_option("float_format", val) 613 | for field in self._field_names: 614 | self._float_format[field] = val 615 | float_format = property(_get_float_format, _set_float_format) 616 | 617 | def _get_padding_width(self): 618 | """The number of empty spaces between a column's edge and its content 619 | 620 | Arguments: 621 | 622 | padding_width - number of spaces, must be a positive integer""" 623 | return self._padding_width 624 | def _set_padding_width(self, val): 625 | self._validate_option("padding_width", val) 626 | self._padding_width = val 627 | padding_width = property(_get_padding_width, _set_padding_width) 628 | 629 | def _get_left_padding_width(self): 630 | """The number of empty spaces between a column's left edge and its content 631 | 632 | Arguments: 633 | 634 | left_padding - number of spaces, must be a positive integer""" 635 | return self._left_padding_width 636 | def _set_left_padding_width(self, val): 637 | self._validate_option("left_padding_width", val) 638 | self._left_padding_width = val 639 | left_padding_width = property(_get_left_padding_width, _set_left_padding_width) 640 | 641 | def _get_right_padding_width(self): 642 | """The number of empty spaces between a column's right edge and its content 643 | 644 | Arguments: 645 | 646 | right_padding - number of spaces, must be a positive integer""" 647 | return self._right_padding_width 648 | def _set_right_padding_width(self, val): 649 | self._validate_option("right_padding_width", val) 650 | self._right_padding_width = val 651 | right_padding_width = property(_get_right_padding_width, _set_right_padding_width) 652 | 653 | def _get_vertical_char(self): 654 | """The charcter used when printing table borders to draw vertical lines 655 | 656 | Arguments: 657 | 658 | vertical_char - single character string used to draw vertical lines""" 659 | return self._vertical_char 660 | def _set_vertical_char(self, val): 661 | val = self._unicode(val) 662 | self._validate_option("vertical_char", val) 663 | self._vertical_char = val 664 | vertical_char = property(_get_vertical_char, _set_vertical_char) 665 | 666 | def _get_horizontal_char(self): 667 | """The charcter used when printing table borders to draw horizontal lines 668 | 669 | Arguments: 670 | 671 | horizontal_char - single character string used to draw horizontal lines""" 672 | return self._horizontal_char 673 | def _set_horizontal_char(self, val): 674 | val = self._unicode(val) 675 | self._validate_option("horizontal_char", val) 676 | self._horizontal_char = val 677 | horizontal_char = property(_get_horizontal_char, _set_horizontal_char) 678 | 679 | def _get_junction_char(self): 680 | """The charcter used when printing table borders to draw line junctions 681 | 682 | Arguments: 683 | 684 | junction_char - single character string used to draw line junctions""" 685 | return self._junction_char 686 | def _set_junction_char(self, val): 687 | val = self._unicode(val) 688 | self._validate_option("vertical_char", val) 689 | self._junction_char = val 690 | junction_char = property(_get_junction_char, _set_junction_char) 691 | 692 | def _get_format(self): 693 | """Controls whether or not HTML tables are formatted to match styling options 694 | 695 | Arguments: 696 | 697 | format - True or False""" 698 | return self._format 699 | def _set_format(self, val): 700 | self._validate_option("format", val) 701 | self._format = val 702 | format = property(_get_format, _set_format) 703 | 704 | def _get_print_empty(self): 705 | """Controls whether or not empty tables produce a header and frame or just an empty string 706 | 707 | Arguments: 708 | 709 | print_empty - True or False""" 710 | return self._print_empty 711 | def _set_print_empty(self, val): 712 | self._validate_option("print_empty", val) 713 | self._print_empty = val 714 | print_empty = property(_get_print_empty, _set_print_empty) 715 | 716 | def _get_attributes(self): 717 | """A dictionary of HTML attribute name/value pairs to be included in the tag when printing HTML 718 | 719 | Arguments: 720 | 721 | attributes - dictionary of attributes""" 722 | return self._attributes 723 | def _set_attributes(self, val): 724 | self._validate_option("attributes", val) 725 | self._attributes = val 726 | attributes = property(_get_attributes, _set_attributes) 727 | 728 | ############################## 729 | # OPTION MIXER # 730 | ############################## 731 | 732 | def _get_options(self, kwargs): 733 | 734 | options = {} 735 | for option in self._options: 736 | if option in kwargs: 737 | self._validate_option(option, kwargs[option]) 738 | options[option] = kwargs[option] 739 | else: 740 | options[option] = getattr(self, "_"+option) 741 | return options 742 | 743 | ############################## 744 | # PRESET STYLE LOGIC # 745 | ############################## 746 | 747 | def set_style(self, style): 748 | 749 | if style == DEFAULT: 750 | self._set_default_style() 751 | elif style == MSWORD_FRIENDLY: 752 | self._set_msword_style() 753 | elif style == PLAIN_COLUMNS: 754 | self._set_columns_style() 755 | elif style == RANDOM: 756 | self._set_random_style() 757 | else: 758 | raise Exception("Invalid pre-set style!") 759 | 760 | def _set_default_style(self): 761 | 762 | self.header = True 763 | self.border = True 764 | self._hrules = FRAME 765 | self._vrules = ALL 766 | self.padding_width = 1 767 | self.left_padding_width = 1 768 | self.right_padding_width = 1 769 | self.vertical_char = "|" 770 | self.horizontal_char = "-" 771 | self.junction_char = "+" 772 | 773 | def _set_msword_style(self): 774 | 775 | self.header = True 776 | self.border = True 777 | self._hrules = NONE 778 | self.padding_width = 1 779 | self.left_padding_width = 1 780 | self.right_padding_width = 1 781 | self.vertical_char = "|" 782 | 783 | def _set_columns_style(self): 784 | 785 | self.header = True 786 | self.border = False 787 | self.padding_width = 1 788 | self.left_padding_width = 0 789 | self.right_padding_width = 8 790 | 791 | def _set_random_style(self): 792 | 793 | # Just for fun! 794 | self.header = random.choice((True, False)) 795 | self.border = random.choice((True, False)) 796 | self._hrules = random.choice((ALL, FRAME, HEADER, NONE)) 797 | self._vrules = random.choice((ALL, FRAME, NONE)) 798 | self.left_padding_width = random.randint(0,5) 799 | self.right_padding_width = random.randint(0,5) 800 | self.vertical_char = random.choice("~!@#$%^&*()_+|-=\{}[];':\",./;<>?") 801 | self.horizontal_char = random.choice("~!@#$%^&*()_+|-=\{}[];':\",./;<>?") 802 | self.junction_char = random.choice("~!@#$%^&*()_+|-=\{}[];':\",./;<>?") 803 | 804 | ############################## 805 | # DATA INPUT METHODS # 806 | ############################## 807 | 808 | def add_row(self, row): 809 | 810 | """Add a row to the table 811 | 812 | Arguments: 813 | 814 | row - row of data, should be a list with as many elements as the table 815 | has fields""" 816 | 817 | if self._field_names and len(row) != len(self._field_names): 818 | raise Exception("Row has incorrect number of values, (actual) %d!=%d (expected)" %(len(row),len(self._field_names))) 819 | if not self._field_names: 820 | self.field_names = [("Field %d" % (n+1)) for n in range(0,len(row))] 821 | self._rows.append(list(row)) 822 | 823 | def del_row(self, row_index): 824 | 825 | """Delete a row to the table 826 | 827 | Arguments: 828 | 829 | row_index - The index of the row you want to delete. Indexing starts at 0.""" 830 | 831 | if row_index > len(self._rows)-1: 832 | raise Exception("Cant delete row at index %d, table only has %d rows!" % (row_index, len(self._rows))) 833 | del self._rows[row_index] 834 | 835 | def add_column(self, fieldname, column, align="c", valign="t"): 836 | 837 | """Add a column to the table. 838 | 839 | Arguments: 840 | 841 | fieldname - name of the field to contain the new column of data 842 | column - column of data, should be a list with as many elements as the 843 | table has rows 844 | align - desired alignment for this column - "l" for left, "c" for centre and "r" for right 845 | valign - desired vertical alignment for new columns - "t" for top, "m" for middle and "b" for bottom""" 846 | 847 | if len(self._rows) in (0, len(column)): 848 | self._validate_align(align) 849 | self._validate_valign(valign) 850 | self._field_names.append(fieldname) 851 | self._align[fieldname] = align 852 | self._valign[fieldname] = valign 853 | for i in range(0, len(column)): 854 | if len(self._rows) < i+1: 855 | self._rows.append([]) 856 | self._rows[i].append(column[i]) 857 | else: 858 | raise Exception("Column length %d does not match number of rows %d!" % (len(column), len(self._rows))) 859 | 860 | def clear_rows(self): 861 | 862 | """Delete all rows from the table but keep the current field names""" 863 | 864 | self._rows = [] 865 | 866 | def clear(self): 867 | 868 | """Delete all rows and field names from the table, maintaining nothing but styling options""" 869 | 870 | self._rows = [] 871 | self._field_names = [] 872 | self._widths = [] 873 | 874 | ############################## 875 | # MISC PUBLIC METHODS # 876 | ############################## 877 | 878 | def copy(self): 879 | return copy.deepcopy(self) 880 | 881 | ############################## 882 | # MISC PRIVATE METHODS # 883 | ############################## 884 | 885 | def _format_value(self, field, value): 886 | if isinstance(value, int) and field in self._int_format: 887 | value = self._unicode(("%%%sd" % self._int_format[field]) % value) 888 | elif isinstance(value, float) and field in self._float_format: 889 | value = self._unicode(("%%%sf" % self._float_format[field]) % value) 890 | return self._unicode(value) 891 | 892 | def _compute_widths(self, rows, options): 893 | if options["header"]: 894 | widths = [_get_size(field)[0] for field in self._field_names] 895 | else: 896 | widths = len(self.field_names) * [0] 897 | for row in rows: 898 | for index, value in enumerate(row): 899 | fieldname = self.field_names[index] 900 | if fieldname in self.max_width: 901 | widths[index] = max(widths[index], min(_get_size(value)[0], self.max_width[fieldname])) 902 | else: 903 | widths[index] = max(widths[index], _get_size(value)[0]) 904 | self._widths = widths 905 | 906 | def _get_padding_widths(self, options): 907 | 908 | if options["left_padding_width"] is not None: 909 | lpad = options["left_padding_width"] 910 | else: 911 | lpad = options["padding_width"] 912 | if options["right_padding_width"] is not None: 913 | rpad = options["right_padding_width"] 914 | else: 915 | rpad = options["padding_width"] 916 | return lpad, rpad 917 | 918 | def _get_rows(self, options): 919 | """Return only those data rows that should be printed, based on slicing and sorting. 920 | 921 | Arguments: 922 | 923 | options - dictionary of option settings.""" 924 | 925 | # Make a copy of only those rows in the slice range 926 | rows = copy.deepcopy(self._rows[options["start"]:options["end"]]) 927 | # Sort if necessary 928 | if options["sortby"]: 929 | sortindex = self._field_names.index(options["sortby"]) 930 | # Decorate 931 | rows = [[row[sortindex]]+row for row in rows] 932 | # Sort 933 | rows.sort(reverse=options["reversesort"], key=options["sort_key"]) 934 | # Undecorate 935 | rows = [row[1:] for row in rows] 936 | return rows 937 | 938 | def _format_row(self, row, options): 939 | return [self._format_value(field, value) for (field, value) in zip(self._field_names, row)] 940 | 941 | def _format_rows(self, rows, options): 942 | return [self._format_row(row, options) for row in rows] 943 | 944 | ############################## 945 | # PLAIN TEXT STRING METHODS # 946 | ############################## 947 | 948 | def get_string(self, **kwargs): 949 | 950 | """Return string representation of table in current state. 951 | 952 | Arguments: 953 | 954 | start - index of first data row to include in output 955 | end - index of last data row to include in output PLUS ONE (list slice style) 956 | fields - names of fields (columns) to include 957 | header - print a header showing field names (True or False) 958 | border - print a border around the table (True or False) 959 | hrules - controls printing of horizontal rules after rows. Allowed values: ALL, FRAME, HEADER, NONE 960 | vrules - controls printing of vertical rules between columns. Allowed values: FRAME, ALL, NONE 961 | int_format - controls formatting of integer data 962 | float_format - controls formatting of floating point data 963 | padding_width - number of spaces on either side of column data (only used if left and right paddings are None) 964 | left_padding_width - number of spaces on left hand side of column data 965 | right_padding_width - number of spaces on right hand side of column data 966 | vertical_char - single character string used to draw vertical lines 967 | horizontal_char - single character string used to draw horizontal lines 968 | junction_char - single character string used to draw line junctions 969 | sortby - name of field to sort rows by 970 | sort_key - sorting key function, applied to data points before sorting 971 | reversesort - True or False to sort in descending or ascending order 972 | print empty - if True, stringify just the header for an empty table, if False return an empty string """ 973 | 974 | options = self._get_options(kwargs) 975 | 976 | lines = [] 977 | 978 | # Don't think too hard about an empty table 979 | # Is this the desired behaviour? Maybe we should still print the header? 980 | if self.rowcount == 0 and (not options["print_empty"] or not options["border"]): 981 | return "" 982 | 983 | # Get the rows we need to print, taking into account slicing, sorting, etc. 984 | rows = self._get_rows(options) 985 | 986 | # Turn all data in all rows into Unicode, formatted as desired 987 | formatted_rows = self._format_rows(rows, options) 988 | 989 | # Compute column widths 990 | self._compute_widths(formatted_rows, options) 991 | 992 | # Add header or top of border 993 | self._hrule = self._stringify_hrule(options) 994 | if options["header"]: 995 | lines.append(self._stringify_header(options)) 996 | elif options["border"] and options["hrules"] in (ALL, FRAME): 997 | lines.append(self._hrule) 998 | 999 | # Add rows 1000 | for row in formatted_rows: 1001 | lines.append(self._stringify_row(row, options)) 1002 | 1003 | # Add bottom of border 1004 | if options["border"] and options["hrules"] == FRAME: 1005 | lines.append(self._hrule) 1006 | 1007 | return self._unicode("\n").join(lines) 1008 | 1009 | def _stringify_hrule(self, options): 1010 | 1011 | if not options["border"]: 1012 | return "" 1013 | lpad, rpad = self._get_padding_widths(options) 1014 | if options['vrules'] in (ALL, FRAME): 1015 | bits = [options["junction_char"]] 1016 | else: 1017 | bits = [options["horizontal_char"]] 1018 | # For tables with no data or fieldnames 1019 | if not self._field_names: 1020 | bits.append(options["junction_char"]) 1021 | return "".join(bits) 1022 | for field, width in zip(self._field_names, self._widths): 1023 | if options["fields"] and field not in options["fields"]: 1024 | continue 1025 | bits.append((width+lpad+rpad)*options["horizontal_char"]) 1026 | if options['vrules'] == ALL: 1027 | bits.append(options["junction_char"]) 1028 | else: 1029 | bits.append(options["horizontal_char"]) 1030 | if options["vrules"] == FRAME: 1031 | bits.pop() 1032 | bits.append(options["junction_char"]) 1033 | return "".join(bits) 1034 | 1035 | def _stringify_header(self, options): 1036 | 1037 | bits = [] 1038 | lpad, rpad = self._get_padding_widths(options) 1039 | if options["border"]: 1040 | if options["hrules"] in (ALL, FRAME): 1041 | bits.append(self._hrule) 1042 | bits.append("\n") 1043 | if options["vrules"] in (ALL, FRAME): 1044 | bits.append(options["vertical_char"]) 1045 | else: 1046 | bits.append(" ") 1047 | # For tables with no data or field names 1048 | if not self._field_names: 1049 | if options["vrules"] in (ALL, FRAME): 1050 | bits.append(options["vertical_char"]) 1051 | else: 1052 | bits.append(" ") 1053 | for field, width, in zip(self._field_names, self._widths): 1054 | if options["fields"] and field not in options["fields"]: 1055 | continue 1056 | if self._header_style == "cap": 1057 | fieldname = field.capitalize() 1058 | elif self._header_style == "title": 1059 | fieldname = field.title() 1060 | elif self._header_style == "upper": 1061 | fieldname = field.upper() 1062 | elif self._header_style == "lower": 1063 | fieldname = field.lower() 1064 | else: 1065 | fieldname = field 1066 | bits.append(" " * lpad + self._justify(fieldname, width, self._align[field]) + " " * rpad) 1067 | if options["border"]: 1068 | if options["vrules"] == ALL: 1069 | bits.append(options["vertical_char"]) 1070 | else: 1071 | bits.append(" ") 1072 | # If vrules is FRAME, then we just appended a space at the end 1073 | # of the last field, when we really want a vertical character 1074 | if options["border"] and options["vrules"] == FRAME: 1075 | bits.pop() 1076 | bits.append(options["vertical_char"]) 1077 | if options["border"] and options["hrules"] != NONE: 1078 | bits.append("\n") 1079 | bits.append(self._hrule) 1080 | return "".join(bits) 1081 | 1082 | def _stringify_row(self, row, options): 1083 | 1084 | for index, field, value, width, in zip(range(0,len(row)), self._field_names, row, self._widths): 1085 | # Enforce max widths 1086 | lines = value.split("\n") 1087 | new_lines = [] 1088 | for line in lines: 1089 | if _str_block_width(line) > width: 1090 | line = textwrap.fill(line, width) 1091 | new_lines.append(line) 1092 | lines = new_lines 1093 | value = "\n".join(lines) 1094 | row[index] = value 1095 | 1096 | row_height = 0 1097 | for c in row: 1098 | h = _get_size(c)[1] 1099 | if h > row_height: 1100 | row_height = h 1101 | 1102 | bits = [] 1103 | lpad, rpad = self._get_padding_widths(options) 1104 | for y in range(0, row_height): 1105 | bits.append([]) 1106 | if options["border"]: 1107 | if options["vrules"] in (ALL, FRAME): 1108 | bits[y].append(self.vertical_char) 1109 | else: 1110 | bits[y].append(" ") 1111 | 1112 | for field, value, width, in zip(self._field_names, row, self._widths): 1113 | 1114 | valign = self._valign[field] 1115 | lines = value.split("\n") 1116 | dHeight = row_height - len(lines) 1117 | if dHeight: 1118 | if valign == "m": 1119 | lines = [""] * int(dHeight / 2) + lines + [""] * (dHeight - int(dHeight / 2)) 1120 | elif valign == "b": 1121 | lines = [""] * dHeight + lines 1122 | else: 1123 | lines = lines + [""] * dHeight 1124 | 1125 | y = 0 1126 | for l in lines: 1127 | if options["fields"] and field not in options["fields"]: 1128 | continue 1129 | 1130 | bits[y].append(" " * lpad + self._justify(l, width, self._align[field]) + " " * rpad) 1131 | if options["border"]: 1132 | if options["vrules"] == ALL: 1133 | bits[y].append(self.vertical_char) 1134 | else: 1135 | bits[y].append(" ") 1136 | y += 1 1137 | 1138 | # If vrules is FRAME, then we just appended a space at the end 1139 | # of the last field, when we really want a vertical character 1140 | for y in range(0, row_height): 1141 | if options["border"] and options["vrules"] == FRAME: 1142 | bits[y].pop() 1143 | bits[y].append(options["vertical_char"]) 1144 | 1145 | if options["border"] and options["hrules"]== ALL: 1146 | bits[row_height-1].append("\n") 1147 | bits[row_height-1].append(self._hrule) 1148 | 1149 | for y in range(0, row_height): 1150 | bits[y] = "".join(bits[y]) 1151 | 1152 | return "\n".join(bits) 1153 | 1154 | ############################## 1155 | # HTML STRING METHODS # 1156 | ############################## 1157 | 1158 | def get_html_string(self, **kwargs): 1159 | 1160 | """Return string representation of HTML formatted version of table in current state. 1161 | 1162 | Arguments: 1163 | 1164 | start - index of first data row to include in output 1165 | end - index of last data row to include in output PLUS ONE (list slice style) 1166 | fields - names of fields (columns) to include 1167 | header - print a header showing field names (True or False) 1168 | border - print a border around the table (True or False) 1169 | hrules - controls printing of horizontal rules after rows. Allowed values: ALL, FRAME, HEADER, NONE 1170 | vrules - controls printing of vertical rules between columns. Allowed values: FRAME, ALL, NONE 1171 | int_format - controls formatting of integer data 1172 | float_format - controls formatting of floating point data 1173 | padding_width - number of spaces on either side of column data (only used if left and right paddings are None) 1174 | left_padding_width - number of spaces on left hand side of column data 1175 | right_padding_width - number of spaces on right hand side of column data 1176 | sortby - name of field to sort rows by 1177 | sort_key - sorting key function, applied to data points before sorting 1178 | attributes - dictionary of name/value pairs to include as HTML attributes in the
tag 1179 | xhtml - print
tags if True,
tags if false""" 1180 | 1181 | options = self._get_options(kwargs) 1182 | 1183 | if options["format"]: 1184 | string = self._get_formatted_html_string(options) 1185 | else: 1186 | string = self._get_simple_html_string(options) 1187 | 1188 | return string 1189 | 1190 | def _get_simple_html_string(self, options): 1191 | 1192 | lines = [] 1193 | if options["xhtml"]: 1194 | linebreak = "
" 1195 | else: 1196 | linebreak = "
" 1197 | 1198 | open_tag = [] 1199 | open_tag.append("") 1204 | lines.append("".join(open_tag)) 1205 | 1206 | # Headers 1207 | if options["header"]: 1208 | lines.append(" ") 1209 | for field in self._field_names: 1210 | if options["fields"] and field not in options["fields"]: 1211 | continue 1212 | lines.append(" " % escape(field).replace("\n", linebreak)) 1213 | lines.append(" ") 1214 | 1215 | # Data 1216 | rows = self._get_rows(options) 1217 | formatted_rows = self._format_rows(rows, options) 1218 | for row in formatted_rows: 1219 | lines.append(" ") 1220 | for field, datum in zip(self._field_names, row): 1221 | if options["fields"] and field not in options["fields"]: 1222 | continue 1223 | lines.append(" " % escape(datum).replace("\n", linebreak)) 1224 | lines.append(" ") 1225 | 1226 | lines.append("
%s
%s
") 1227 | 1228 | return self._unicode("\n").join(lines) 1229 | 1230 | def _get_formatted_html_string(self, options): 1231 | 1232 | lines = [] 1233 | lpad, rpad = self._get_padding_widths(options) 1234 | if options["xhtml"]: 1235 | linebreak = "
" 1236 | else: 1237 | linebreak = "
" 1238 | 1239 | open_tag = [] 1240 | open_tag.append("") 1260 | lines.append("".join(open_tag)) 1261 | 1262 | # Headers 1263 | if options["header"]: 1264 | lines.append(" ") 1265 | for field in self._field_names: 1266 | if options["fields"] and field not in options["fields"]: 1267 | continue 1268 | lines.append(" %s" % (lpad, rpad, escape(field).replace("\n", linebreak))) 1269 | lines.append(" ") 1270 | 1271 | # Data 1272 | rows = self._get_rows(options) 1273 | formatted_rows = self._format_rows(rows, options) 1274 | aligns = [] 1275 | valigns = [] 1276 | for field in self._field_names: 1277 | aligns.append({ "l" : "left", "r" : "right", "c" : "center" }[self._align[field]]) 1278 | valigns.append({"t" : "top", "m" : "middle", "b" : "bottom"}[self._valign[field]]) 1279 | for row in formatted_rows: 1280 | lines.append(" ") 1281 | for field, datum, align, valign in zip(self._field_names, row, aligns, valigns): 1282 | if options["fields"] and field not in options["fields"]: 1283 | continue 1284 | lines.append(" %s" % (lpad, rpad, align, valign, escape(datum).replace("\n", linebreak))) 1285 | lines.append(" ") 1286 | lines.append("") 1287 | 1288 | return self._unicode("\n").join(lines) 1289 | 1290 | ############################## 1291 | # UNICODE WIDTH FUNCTIONS # 1292 | ############################## 1293 | 1294 | def _char_block_width(char): 1295 | # Basic Latin, which is probably the most common case 1296 | #if char in xrange(0x0021, 0x007e): 1297 | #if char >= 0x0021 and char <= 0x007e: 1298 | if 0x0021 <= char <= 0x007e: 1299 | return 1 1300 | # Chinese, Japanese, Korean (common) 1301 | if 0x4e00 <= char <= 0x9fff: 1302 | return 2 1303 | # Hangul 1304 | if 0xac00 <= char <= 0xd7af: 1305 | return 2 1306 | # Combining? 1307 | if unicodedata.combining(uni_chr(char)): 1308 | return 0 1309 | # Hiragana and Katakana 1310 | if 0x3040 <= char <= 0x309f or 0x30a0 <= char <= 0x30ff: 1311 | return 2 1312 | # Full-width Latin characters 1313 | if 0xff01 <= char <= 0xff60: 1314 | return 2 1315 | # CJK punctuation 1316 | if 0x3000 <= char <= 0x303e: 1317 | return 2 1318 | # Backspace and delete 1319 | if char in (0x0008, 0x007f): 1320 | return -1 1321 | # Other control characters 1322 | elif char in (0x0000, 0x001f): 1323 | return 0 1324 | # Take a guess 1325 | return 1 1326 | 1327 | def _str_block_width(val): 1328 | 1329 | return sum(itermap(_char_block_width, itermap(ord, _re.sub("", val)))) 1330 | 1331 | ############################## 1332 | # TABLE FACTORIES # 1333 | ############################## 1334 | 1335 | def from_csv(fp, field_names = None, **kwargs): 1336 | 1337 | dialect = csv.Sniffer().sniff(fp.read(1024)) 1338 | fp.seek(0) 1339 | reader = csv.reader(fp, dialect) 1340 | 1341 | table = PrettyTable(**kwargs) 1342 | if field_names: 1343 | table.field_names = field_names 1344 | else: 1345 | if py3k: 1346 | table.field_names = [x.strip() for x in next(reader)] 1347 | else: 1348 | table.field_names = [x.strip() for x in reader.next()] 1349 | 1350 | for row in reader: 1351 | table.add_row([x.strip() for x in row]) 1352 | 1353 | return table 1354 | 1355 | def from_db_cursor(cursor, **kwargs): 1356 | 1357 | if cursor.description: 1358 | table = PrettyTable(**kwargs) 1359 | table.field_names = [col[0] for col in cursor.description] 1360 | for row in cursor.fetchall(): 1361 | table.add_row(row) 1362 | return table 1363 | 1364 | class TableHandler(HTMLParser): 1365 | 1366 | def __init__(self, **kwargs): 1367 | HTMLParser.__init__(self) 1368 | self.kwargs = kwargs 1369 | self.tables = [] 1370 | self.last_row = [] 1371 | self.rows = [] 1372 | self.max_row_width = 0 1373 | self.active = None 1374 | self.last_content = "" 1375 | self.is_last_row_header = False 1376 | 1377 | def handle_starttag(self,tag, attrs): 1378 | self.active = tag 1379 | if tag == "th": 1380 | self.is_last_row_header = True 1381 | 1382 | def handle_endtag(self,tag): 1383 | if tag in ["th", "td"]: 1384 | stripped_content = self.last_content.strip() 1385 | self.last_row.append(stripped_content) 1386 | if tag == "tr": 1387 | self.rows.append( 1388 | (self.last_row, self.is_last_row_header)) 1389 | self.max_row_width = max(self.max_row_width, len(self.last_row)) 1390 | self.last_row = [] 1391 | self.is_last_row_header = False 1392 | if tag == "table": 1393 | table = self.generate_table(self.rows) 1394 | self.tables.append(table) 1395 | self.rows = [] 1396 | self.last_content = " " 1397 | self.active = None 1398 | 1399 | 1400 | def handle_data(self, data): 1401 | self.last_content += data 1402 | 1403 | def generate_table(self, rows): 1404 | """ 1405 | Generates from a list of rows a PrettyTable object. 1406 | """ 1407 | table = PrettyTable(**self.kwargs) 1408 | for row in self.rows: 1409 | if len(row[0]) < self.max_row_width: 1410 | appends = self.max_row_width - len(row[0]) 1411 | for i in range(1,appends): 1412 | row[0].append("-") 1413 | 1414 | if row[1] == True: 1415 | self.make_fields_unique(row[0]) 1416 | table.field_names = row[0] 1417 | else: 1418 | table.add_row(row[0]) 1419 | return table 1420 | 1421 | def make_fields_unique(self, fields): 1422 | """ 1423 | iterates over the row and make each field unique 1424 | """ 1425 | for i in range(0, len(fields)): 1426 | for j in range(i+1, len(fields)): 1427 | if fields[i] == fields[j]: 1428 | fields[j] += "'" 1429 | 1430 | def from_html(html_code, **kwargs): 1431 | """ 1432 | Generates a list of PrettyTables from a string of HTML code. Each in 1433 | the HTML becomes one PrettyTable object. 1434 | """ 1435 | 1436 | parser = TableHandler(**kwargs) 1437 | parser.feed(html_code) 1438 | return parser.tables 1439 | 1440 | def from_html_one(html_code, **kwargs): 1441 | """ 1442 | Generates a PrettyTables from a string of HTML code which contains only a 1443 | single
1444 | """ 1445 | 1446 | tables = from_html(html_code, **kwargs) 1447 | try: 1448 | assert len(tables) == 1 1449 | except AssertionError: 1450 | raise Exception("More than one
in provided HTML code! Use from_html instead.") 1451 | return tables[0] 1452 | 1453 | ############################## 1454 | # MAIN (TEST FUNCTION) # 1455 | ############################## 1456 | 1457 | def main(): 1458 | 1459 | x = PrettyTable(["City name", "Area", "Population", "Annual Rainfall"]) 1460 | x.sortby = "Population" 1461 | x.reversesort = True 1462 | x.int_format["Area"] = "04d" 1463 | x.float_format = "6.1f" 1464 | x.align["City name"] = "l" # Left align city names 1465 | x.add_row(["Adelaide", 1295, 1158259, 600.5]) 1466 | x.add_row(["Brisbane", 5905, 1857594, 1146.4]) 1467 | x.add_row(["Darwin", 112, 120900, 1714.7]) 1468 | x.add_row(["Hobart", 1357, 205556, 619.5]) 1469 | x.add_row(["Sydney", 2058, 4336374, 1214.8]) 1470 | x.add_row(["Melbourne", 1566, 3806092, 646.9]) 1471 | x.add_row(["Perth", 5386, 1554769, 869.4]) 1472 | print(x) 1473 | 1474 | if __name__ == "__main__": 1475 | main() 1476 | -------------------------------------------------------------------------------- /prettytable.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/giantbranch/mipsAudit/a8e5576461a75c92f4ac7f8ecaa9ee096b8db371/prettytable.pyc --------------------------------------------------------------------------------