├── preview.gif ├── README.md ├── LICENSE ├── .gitignore ├── downloader.py └── thirdpart ├── termcolor.py └── prettytable.py /preview.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/boy-hack/vulhub-downloader/HEAD/preview.gif -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # vulhub-downloader 2 | vulhub下载器，可则需下载对应环境(不需要下载整个环境) 3 | - 基于python3.x编写 4 | - 不需要下载整个环境 5 | - 可搜索 app、name、cve对应下载 6 | 7 | ## 待完成 8 | 9 | - [x] 完善异常处理 10 | - [x] 支持关键词正则搜索 11 | - [x] 界面美化 12 | - [ ] 支持使用一句话运行此脚本 13 | - [ ] 支持命令行参数 14 | - [ ] 支持python2脚本 15 | 16 | ## 演示 17 | 18 | ![preview](preview.gif) 19 | 20 | ## Thx 21 | 22 | [https://github.com/vulhub/vulhub](https://github.com/vulhub/vulhub) 23 | 24 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 boyhack 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | db.sqlite3 58 | 59 | # Flask stuff: 60 | instance/ 61 | .webassets-cache 62 | 63 | # Scrapy stuff: 64 | .scrapy 65 | 66 | # Sphinx documentation 67 | docs/_build/ 68 | 69 | # PyBuilder 70 | target/ 71 | 72 | # Jupyter Notebook 73 | .ipynb_checkpoints 74 | 75 | # pyenv 76 | .python-version 77 | 78 | # celery beat schedule file 79 | celerybeat-schedule 80 | 81 | # SageMath parsed files 82 | *.sage.py 83 | 84 | # Environments 85 | .env 86 | .venv 87 | env/ 88 | venv/ 89 | ENV/ 90 | env.bak/ 91 | venv.bak/ 92 | 93 | # Spyder project settings 94 | .spyderproject 95 | .spyproject 96 | 97 | # Rope project settings 98 | .ropeproject 99 | 100 | # mkdocs documentation 101 | /site 102 | 103 | # mypy 104 | .mypy_cache/ 105 | .idea/ 106 | .vscode/ 107 | .DS_Store -------------------------------------------------------------------------------- /downloader.py: -------------------------------------------------------------------------------- 1 | from thirdpart.prettytable import PrettyTable 2 | import sys 3 | import time 4 | import json 5 | import urllib.request 6 | from thirdpart.termcolor import cprint 7 | import os 8 | import re 9 | 10 | 11 | def wget(url): 12 | header = { 13 | 'User-Agent': 'Mozilla/5.0 (X11; Fedora; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.36' 14 | } 15 | request = urllib.request.Request(url, headers=header) 16 | reponse = urllib.request.urlopen(request).read() 17 | return reponse 18 | 19 | 20 | class Vulhub_downloader(object): 21 | 22 | def __init__(self): 23 | self.origin = "https://raw.githubusercontent.com/vulhub/vulhub-org/master/src/environments.json" 24 | self.api = "https://api.github.com/repos/vulhub/vulhub/contents/" 25 | self.directory = "" # 初始化目录 26 | 27 | self.originText = json.loads(wget(self.origin).decode("utf8")) 28 | 29 | def parse_github(self, path): 30 | url = self.api + path 31 | apiText = json.loads(wget(url).decode("utf8")) 32 | return apiText 33 | 34 | def download(self, directory, data): 35 | for item in data: 36 | name = item.get("name") 37 | app = item.get("app") 38 | cve = item.get("cve") 39 | path = item.get("path") 40 | api = self.parse_github(path) 41 | download_directory = os.path.join(directory, path) 42 | total = len(api) 43 | index = 0 44 | _except = False 45 | for temp in api: 46 | name = temp.get("name") 47 | download_url = temp.get("download_url") 48 | 49 | type = temp.get("type") 50 | if type == "file": 51 | print("download:{}".format(name)) 52 | try: 53 | content = wget(download_url) 54 | except: 55 | _except = True 56 | break 57 | if not os.path.exists(download_directory): 58 | os.makedirs(download_directory) 59 | with open(os.path.join(download_directory, name), 'wb') as f: 60 | f.write(content) 61 | 62 | index += 1 63 | if _except is False: 64 | cprint("success","green") 65 | x = PrettyTable(["name", "app", "cve", "directory"]) 66 | x.align["name"] = "l" # 以name字段左对齐 67 | x.padding_width = 1 # 填充宽度 68 | x.add_row([name, app, str(cve), download_directory]) 69 | print(x) 70 | else: 71 | cprint("下载失败","red") 72 | 73 | def search(self, keywords): 74 | ''' 75 | 支持关键词搜索(可搜索app、name、cve)，可使用g:进行正则搜索，返回搜索到的数据 76 | :param keywords: 77 | :return: 78 | ''' 79 | result = [] 80 | for item in self.originText: 81 | name = item.get("name") 82 | app = item.get("app") 83 | cve = item.get("cve") 84 | if cve is None: 85 | cve = "" 86 | path = item.get("path") 87 | keylower = keywords.lower() 88 | if keylower.startswith("g:"): 89 | keylower = keylower[2:] 90 | if re.search(keylower,name) or re.search(keylower,app) or re.search(keylower,cve): 91 | result.append(item) 92 | else: 93 | if keylower in name.lower() or keylower in app.lower() or keylower in cve.lower(): 94 | result.append(item) 95 | return result 96 | 97 | 98 | def gui(): 99 | down = Vulhub_downloader() 100 | banner = r''' 101 | ❤️ ( ⚫︎ー⚫︎ ) balalala~ 102 | 　／　　　 ＼ 103 | /　　　 ○ 　\ 104 | /　 / 　 ヽ \ 105 | |　/　 　　　 \　| 106 | \Ԏ　 Vulhub-downloader 107 | 　卜−　　 ―イ 108 | 　 \　　/\　 / 109 | 　　 ︶　 ︶ 110 | ''' 111 | print(banner) 112 | print("已加载漏洞环境:{}".format(len(down.originText))) 113 | print("请输入关键词搜索(可搜索app、name、cve,不区分大小写,g:开头使用正则搜索)") 114 | k = input("> ") 115 | data = down.search(k) 116 | 117 | if len(data) == 0: 118 | exit() 119 | 120 | x = PrettyTable(["id", "name", "app", "cve"]) 121 | x.align["name"] = "l" # 以name字段左对齐 122 | x.padding_width = 1 # 填充宽度 123 | index = 1 124 | for item in data: 125 | x.add_row([index, item["name"], item["app"], item["cve"]]) 126 | index += 1 127 | print(x) 128 | print("已搜索到{}个环境".format(len(data))) 129 | print("输入欲下载的id号（多个可用,分割）") 130 | id = input("> ") 131 | ids = id.split(",") 132 | downdata = [] 133 | for item in ids: 134 | item_id = int(item) 135 | downdata.append(data[item_id - 1]) 136 | # 下载目录:（有默认目录） 137 | directory = os.path.join(os.getcwd(), "vulhub") 138 | print("下载目录:(默认:{})".format(directory)) 139 | directory = input("> ") 140 | if directory == "": 141 | directory = os.path.join(os.getcwd(), "vulhub") 142 | # 进度条下载 143 | down.download(directory, downdata) 144 | # 下载完成[保存目录 绿色] /下载失败[原因 红色] 145 | 146 | 147 | if __name__ == '__main__': 148 | try: 149 | gui() 150 | except KeyboardInterrupt: 151 | cprint("User Quit","red") 152 | -------------------------------------------------------------------------------- /thirdpart/termcolor.py: -------------------------------------------------------------------------------- 1 | # coding: utf-8 2 | # Copyright (c) 2008-2011 Volvox Development Team 3 | # 4 | # Permission is hereby granted, free of charge, to any person obtaining a copy 5 | # of this software and associated documentation files (the "Software"), to deal 6 | # in the Software without restriction, including without limitation the rights 7 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 8 | # copies of the Software, and to permit persons to whom the Software is 9 | # furnished to do so, subject to the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be included in 12 | # all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 17 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 19 | # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 20 | # THE SOFTWARE. 21 | # 22 | # Author: Konstantin Lepa 23 | 24 | """ANSII Color formatting for output in terminal.""" 25 | 26 | from __future__ import print_function 27 | 28 | import os 29 | 30 | __ALL__ = [ 'colored', 'cprint' ] 31 | 32 | VERSION = (1, 1, 0) 33 | 34 | ATTRIBUTES = dict( 35 | list(zip([ 36 | 'bold', 37 | 'dark', 38 | '', 39 | 'underline', 40 | 'blink', 41 | '', 42 | 'reverse', 43 | 'concealed' 44 | ], 45 | list(range(1, 9)) 46 | )) 47 | ) 48 | del ATTRIBUTES[''] 49 | 50 | 51 | HIGHLIGHTS = dict( 52 | list(zip([ 53 | 'on_grey', 54 | 'on_red', 55 | 'on_green', 56 | 'on_yellow', 57 | 'on_blue', 58 | 'on_magenta', 59 | 'on_cyan', 60 | 'on_white' 61 | ], 62 | list(range(40, 48)) 63 | )) 64 | ) 65 | 66 | 67 | COLORS = dict( 68 | list(zip([ 69 | 'grey', 70 | 'red', 71 | 'green', 72 | 'yellow', 73 | 'blue', 74 | 'magenta', 75 | 'cyan', 76 | 'white', 77 | ], 78 | list(range(30, 38)) 79 | )) 80 | ) 81 | 82 | 83 | RESET = '\033[0m' 84 | 85 | 86 | def colored(text, color=None, on_color=None, attrs=None): 87 | """Colorize text. 88 | 89 | Available text colors: 90 | red, green, yellow, blue, magenta, cyan, white. 91 | 92 | Available text highlights: 93 | on_red, on_green, on_yellow, on_blue, on_magenta, on_cyan, on_white. 94 | 95 | Available attributes: 96 | bold, dark, underline, blink, reverse, concealed. 97 | 98 | Example: 99 | colored('Hello, World!', 'red', 'on_grey', ['blue', 'blink']) 100 | colored('Hello, World!', 'green') 101 | """ 102 | if os.getenv('ANSI_COLORS_DISABLED') is None: 103 | fmt_str = '\033[%dm%s' 104 | if color is not None: 105 | text = fmt_str % (COLORS[color], text) 106 | 107 | if on_color is not None: 108 | text = fmt_str % (HIGHLIGHTS[on_color], text) 109 | 110 | if attrs is not None: 111 | for attr in attrs: 112 | text = fmt_str % (ATTRIBUTES[attr], text) 113 | 114 | text += RESET 115 | return text 116 | 117 | 118 | def cprint(text, color=None, on_color=None, attrs=None, **kwargs): 119 | """Print colorize text. 120 | 121 | It accepts arguments of print function. 122 | """ 123 | 124 | print((colored(text, color, on_color, attrs)), **kwargs) 125 | 126 | 127 | if __name__ == '__main__': 128 | print('Current terminal type: %s' % os.getenv('TERM')) 129 | print('Test basic colors:') 130 | cprint('Grey color', 'grey') 131 | cprint('Red color', 'red') 132 | cprint('Green color', 'green') 133 | cprint('Yellow color', 'yellow') 134 | cprint('Blue color', 'blue') 135 | cprint('Magenta color', 'magenta') 136 | cprint('Cyan color', 'cyan') 137 | cprint('White color', 'white') 138 | print(('-' * 78)) 139 | 140 | print('Test highlights:') 141 | cprint('On grey color', on_color='on_grey') 142 | cprint('On red color', on_color='on_red') 143 | cprint('On green color', on_color='on_green') 144 | cprint('On yellow color', on_color='on_yellow') 145 | cprint('On blue color', on_color='on_blue') 146 | cprint('On magenta color', on_color='on_magenta') 147 | cprint('On cyan color', on_color='on_cyan') 148 | cprint('On white color', color='grey', on_color='on_white') 149 | print('-' * 78) 150 | 151 | print('Test attributes:') 152 | cprint('Bold grey color', 'grey', attrs=['bold']) 153 | cprint('Dark red color', 'red', attrs=['dark']) 154 | cprint('Underline green color', 'green', attrs=['underline']) 155 | cprint('Blink yellow color', 'yellow', attrs=['blink']) 156 | cprint('Reversed blue color', 'blue', attrs=['reverse']) 157 | cprint('Concealed Magenta color', 'magenta', attrs=['concealed']) 158 | cprint('Bold underline reverse cyan color', 'cyan', 159 | attrs=['bold', 'underline', 'reverse']) 160 | cprint('Dark blink concealed white color', 'white', 161 | attrs=['dark', 'blink', 'concealed']) 162 | print(('-' * 78)) 163 | 164 | print('Test mixing:') 165 | cprint('Underline red on grey color', 'red', 'on_grey', 166 | ['underline']) 167 | cprint('Reversed green on red color', 'green', 'on_red', ['reverse']) 168 | -------------------------------------------------------------------------------- /thirdpart/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 | 77 | def _get_size(text): 78 | lines = text.split("\n") 79 | height = len(lines) 80 | width = max([_str_block_width(line) for line in lines]) 81 | return (width, height) 82 | 83 | 84 | class PrettyTable(object): 85 | 86 | def __init__(self, field_names=None, **kwargs): 87 | 88 | """Return a new PrettyTable instance 89 | 90 | Arguments: 91 | 92 | encoding - Unicode encoding scheme used to decode any encoded input 93 | field_names - list or tuple of field names 94 | fields - list or tuple of field names to include in displays 95 | start - index of first data row to include in output 96 | end - index of last data row to include in output PLUS ONE (list slice style) 97 | header - print a header showing field names (True or False) 98 | header_style - stylisation to apply to field names in header ("cap", "title", "upper", "lower" or None) 99 | border - print a border around the table (True or False) 100 | hrules - controls printing of horizontal rules after rows. Allowed values: FRAME, HEADER, ALL, NONE 101 | vrules - controls printing of vertical rules between columns. Allowed values: FRAME, ALL, NONE 102 | int_format - controls formatting of integer data 103 | float_format - controls formatting of floating point data 104 | padding_width - number of spaces on either side of column data (only used if left and right paddings are None) 105 | left_padding_width - number of spaces on left hand side of column data 106 | right_padding_width - number of spaces on right hand side of column data 107 | vertical_char - single character string used to draw vertical lines 108 | horizontal_char - single character string used to draw horizontal lines 109 | junction_char - single character string used to draw line junctions 110 | sortby - name of field to sort rows by 111 | sort_key - sorting key function, applied to data points before sorting 112 | valign - default valign for each row (None, "t", "m" or "b") 113 | reversesort - True or False to sort in descending or ascending order""" 114 | 115 | self.encoding = kwargs.get("encoding", "UTF-8") 116 | 117 | # Data 118 | self._field_names = [] 119 | self._align = {} 120 | self._valign = {} 121 | self._max_width = {} 122 | self._rows = [] 123 | if field_names: 124 | self.field_names = field_names 125 | else: 126 | self._widths = [] 127 | 128 | # Options 129 | self._options = "start end fields header border sortby reversesort sort_key attributes format hrules vrules".split() 130 | self._options.extend("int_format float_format padding_width left_padding_width right_padding_width".split()) 131 | self._options.extend( 132 | "vertical_char horizontal_char junction_char header_style valign xhtml print_empty".split()) 133 | for option in self._options: 134 | if option in kwargs: 135 | self._validate_option(option, kwargs[option]) 136 | else: 137 | kwargs[option] = None 138 | 139 | self._start = kwargs["start"] or 0 140 | self._end = kwargs["end"] or None 141 | self._fields = kwargs["fields"] or None 142 | 143 | if kwargs["header"] in (True, False): 144 | self._header = kwargs["header"] 145 | else: 146 | self._header = True 147 | self._header_style = kwargs["header_style"] or None 148 | if kwargs["border"] in (True, False): 149 | self._border = kwargs["border"] 150 | else: 151 | self._border = True 152 | self._hrules = kwargs["hrules"] or FRAME 153 | self._vrules = kwargs["vrules"] or ALL 154 | 155 | self._sortby = kwargs["sortby"] or None 156 | if kwargs["reversesort"] in (True, False): 157 | self._reversesort = kwargs["reversesort"] 158 | else: 159 | self._reversesort = False 160 | self._sort_key = kwargs["sort_key"] or (lambda x: x) 161 | 162 | self._int_format = kwargs["int_format"] or {} 163 | self._float_format = kwargs["float_format"] or {} 164 | self._padding_width = kwargs["padding_width"] or 1 165 | self._left_padding_width = kwargs["left_padding_width"] or None 166 | self._right_padding_width = kwargs["right_padding_width"] or None 167 | 168 | self._vertical_char = kwargs["vertical_char"] or self._unicode("|") 169 | self._horizontal_char = kwargs["horizontal_char"] or self._unicode("-") 170 | self._junction_char = kwargs["junction_char"] or self._unicode("+") 171 | 172 | if kwargs["print_empty"] in (True, False): 173 | self._print_empty = kwargs["print_empty"] 174 | else: 175 | self._print_empty = True 176 | self._format = kwargs["format"] or False 177 | self._xhtml = kwargs["xhtml"] or False 178 | self._attributes = kwargs["attributes"] or {} 179 | 180 | def _unicode(self, value): 181 | if not isinstance(value, basestring): 182 | value = str(value) 183 | if not isinstance(value, unicode): 184 | value = unicode(value, self.encoding, "strict") 185 | return value 186 | 187 | def _justify(self, text, width, align): 188 | excess = width - _str_block_width(text) 189 | if align == "l": 190 | return text + excess * " " 191 | elif align == "r": 192 | return excess * " " + text 193 | else: 194 | if excess % 2: 195 | # Uneven padding 196 | # Put more space on right if text is of odd length... 197 | if _str_block_width(text) % 2: 198 | return (excess // 2) * " " + text + (excess // 2 + 1) * " " 199 | # and more space on left if text is of even length 200 | else: 201 | return (excess // 2 + 1) * " " + text + (excess // 2) * " " 202 | # Why distribute extra space this way? To match the behaviour of 203 | # the inbuilt str.center() method. 204 | else: 205 | # Equal padding on either side 206 | return (excess // 2) * " " + text + (excess // 2) * " " 207 | 208 | def __getattr__(self, name): 209 | 210 | if name == "rowcount": 211 | return len(self._rows) 212 | elif name == "colcount": 213 | if self._field_names: 214 | return len(self._field_names) 215 | elif self._rows: 216 | return len(self._rows[0]) 217 | else: 218 | return 0 219 | else: 220 | raise AttributeError(name) 221 | 222 | def __getitem__(self, index): 223 | 224 | new = PrettyTable() 225 | new.field_names = self.field_names 226 | for attr in self._options: 227 | setattr(new, "_" + attr, getattr(self, "_" + attr)) 228 | setattr(new, "_align", getattr(self, "_align")) 229 | if isinstance(index, slice): 230 | for row in self._rows[index]: 231 | new.add_row(row) 232 | elif isinstance(index, int): 233 | new.add_row(self._rows[index]) 234 | else: 235 | raise Exception("Index %s is invalid, must be an integer or slice" % str(index)) 236 | return new 237 | 238 | if py3k: 239 | def __str__(self): 240 | return self.__unicode__() 241 | else: 242 | def __str__(self): 243 | return self.__unicode__().encode(self.encoding) 244 | 245 | def __unicode__(self): 246 | return self.get_string() 247 | 248 | ############################## 249 | # ATTRIBUTE VALIDATORS # 250 | ############################## 251 | 252 | # The method _validate_option is all that should be used elsewhere in the code base to validate options. 253 | # It will call the appropriate validation method for that option. The individual validation methods should 254 | # never need to be called directly (although nothing bad will happen if they *are*). 255 | # Validation happens in TWO places. 256 | # Firstly, in the property setters defined in the ATTRIBUTE MANAGMENT section. 257 | # Secondly, in the _get_options method, where keyword arguments are mixed with persistent settings 258 | 259 | def _validate_option(self, option, val): 260 | if option in ("field_names"): 261 | self._validate_field_names(val) 262 | elif option in ( 263 | "start", "end", "max_width", "padding_width", "left_padding_width", "right_padding_width", "format"): 264 | self._validate_nonnegative_int(option, val) 265 | elif option in ("sortby"): 266 | self._validate_field_name(option, val) 267 | elif option in ("sort_key"): 268 | self._validate_function(option, val) 269 | elif option in ("hrules"): 270 | self._validate_hrules(option, val) 271 | elif option in ("vrules"): 272 | self._validate_vrules(option, val) 273 | elif option in ("fields"): 274 | self._validate_all_field_names(option, val) 275 | elif option in ("header", "border", "reversesort", "xhtml", "print_empty"): 276 | self._validate_true_or_false(option, val) 277 | elif option in ("header_style"): 278 | self._validate_header_style(val) 279 | elif option in ("int_format"): 280 | self._validate_int_format(option, val) 281 | elif option in ("float_format"): 282 | self._validate_float_format(option, val) 283 | elif option in ("vertical_char", "horizontal_char", "junction_char"): 284 | self._validate_single_char(option, val) 285 | elif option in ("attributes"): 286 | self._validate_attributes(option, val) 287 | else: 288 | raise Exception("Unrecognised option: %s!" % option) 289 | 290 | def _validate_field_names(self, val): 291 | # Check for appropriate length 292 | if self._field_names: 293 | try: 294 | assert len(val) == len(self._field_names) 295 | except AssertionError: 296 | raise Exception("Field name list has incorrect number of values, (actual) %d!=%d (expected)" % ( 297 | len(val), len(self._field_names))) 298 | if self._rows: 299 | try: 300 | assert len(val) == len(self._rows[0]) 301 | except AssertionError: 302 | raise Exception("Field name list has incorrect number of values, (actual) %d!=%d (expected)" % ( 303 | len(val), len(self._rows[0]))) 304 | # Check for uniqueness 305 | try: 306 | assert len(val) == len(set(val)) 307 | except AssertionError: 308 | raise Exception("Field names must be unique!") 309 | 310 | def _validate_header_style(self, val): 311 | try: 312 | assert val in ("cap", "title", "upper", "lower", None) 313 | except AssertionError: 314 | raise Exception("Invalid header style, use cap, title, upper, lower or None!") 315 | 316 | def _validate_align(self, val): 317 | try: 318 | assert val in ["l", "c", "r"] 319 | except AssertionError: 320 | raise Exception("Alignment %s is invalid, use l, c or r!" % val) 321 | 322 | def _validate_valign(self, val): 323 | try: 324 | assert val in ["t", "m", "b", None] 325 | except AssertionError: 326 | raise Exception("Alignment %s is invalid, use t, m, b or None!" % val) 327 | 328 | def _validate_nonnegative_int(self, name, val): 329 | try: 330 | assert int(val) >= 0 331 | except AssertionError: 332 | raise Exception("Invalid value for %s: %s!" % (name, self._unicode(val))) 333 | 334 | def _validate_true_or_false(self, name, val): 335 | try: 336 | assert val in (True, False) 337 | except AssertionError: 338 | raise Exception("Invalid value for %s! Must be True or False." % name) 339 | 340 | def _validate_int_format(self, name, val): 341 | if val == "": 342 | return 343 | try: 344 | assert type(val) in (str, unicode) 345 | assert val.isdigit() 346 | except AssertionError: 347 | raise Exception("Invalid value for %s! Must be an integer format string." % name) 348 | 349 | def _validate_float_format(self, name, val): 350 | if val == "": 351 | return 352 | try: 353 | assert type(val) in (str, unicode) 354 | assert "." in val 355 | bits = val.split(".") 356 | assert len(bits) <= 2 357 | assert bits[0] == "" or bits[0].isdigit() 358 | assert bits[1] == "" or bits[1].isdigit() 359 | except AssertionError: 360 | raise Exception("Invalid value for %s! Must be a float format string." % name) 361 | 362 | def _validate_function(self, name, val): 363 | try: 364 | assert hasattr(val, "__call__") 365 | except AssertionError: 366 | raise Exception("Invalid value for %s! Must be a function." % name) 367 | 368 | def _validate_hrules(self, name, val): 369 | try: 370 | assert val in (ALL, FRAME, HEADER, NONE) 371 | except AssertionError: 372 | raise Exception("Invalid value for %s! Must be ALL, FRAME, HEADER or NONE." % name) 373 | 374 | def _validate_vrules(self, name, val): 375 | try: 376 | assert val in (ALL, FRAME, NONE) 377 | except AssertionError: 378 | raise Exception("Invalid value for %s! Must be ALL, FRAME, or NONE." % name) 379 | 380 | def _validate_field_name(self, name, val): 381 | try: 382 | assert (val in self._field_names) or (val is None) 383 | except AssertionError: 384 | raise Exception("Invalid field name: %s!" % val) 385 | 386 | def _validate_all_field_names(self, name, val): 387 | try: 388 | for x in val: 389 | self._validate_field_name(name, x) 390 | except AssertionError: 391 | raise Exception("fields must be a sequence of field names!") 392 | 393 | def _validate_single_char(self, name, val): 394 | try: 395 | assert _str_block_width(val) == 1 396 | except AssertionError: 397 | raise Exception("Invalid value for %s! Must be a string of length 1." % name) 398 | 399 | def _validate_attributes(self, name, val): 400 | try: 401 | assert isinstance(val, dict) 402 | except AssertionError: 403 | raise Exception("attributes must be a dictionary of name/value pairs!") 404 | 405 | ############################## 406 | # ATTRIBUTE MANAGEMENT # 407 | ############################## 408 | 409 | def _get_field_names(self): 410 | return self._field_names 411 | """The names of the fields 412 | 413 | Arguments: 414 | 415 | fields - list or tuple of field names""" 416 | 417 | def _set_field_names(self, val): 418 | val = [self._unicode(x) for x in val] 419 | self._validate_option("field_names", val) 420 | if self._field_names: 421 | old_names = self._field_names[:] 422 | self._field_names = val 423 | if self._align and old_names: 424 | for old_name, new_name in zip(old_names, val): 425 | self._align[new_name] = self._align[old_name] 426 | for old_name in old_names: 427 | if old_name not in self._align: 428 | self._align.pop(old_name) 429 | else: 430 | for field in self._field_names: 431 | self._align[field] = "c" 432 | if self._valign and old_names: 433 | for old_name, new_name in zip(old_names, val): 434 | self._valign[new_name] = self._valign[old_name] 435 | for old_name in old_names: 436 | if old_name not in self._valign: 437 | self._valign.pop(old_name) 438 | else: 439 | for field in self._field_names: 440 | self._valign[field] = "t" 441 | 442 | field_names = property(_get_field_names, _set_field_names) 443 | 444 | def _get_align(self): 445 | return self._align 446 | 447 | def _set_align(self, val): 448 | self._validate_align(val) 449 | for field in self._field_names: 450 | self._align[field] = val 451 | 452 | align = property(_get_align, _set_align) 453 | 454 | def _get_valign(self): 455 | return self._valign 456 | 457 | def _set_valign(self, val): 458 | self._validate_valign(val) 459 | for field in self._field_names: 460 | self._valign[field] = val 461 | 462 | valign = property(_get_valign, _set_valign) 463 | 464 | def _get_max_width(self): 465 | return self._max_width 466 | 467 | def _set_max_width(self, val): 468 | self._validate_option("max_width", val) 469 | for field in self._field_names: 470 | self._max_width[field] = val 471 | 472 | max_width = property(_get_max_width, _set_max_width) 473 | 474 | def _get_fields(self): 475 | """List or tuple of field names to include in displays 476 | 477 | Arguments: 478 | 479 | fields - list or tuple of field names to include in displays""" 480 | return self._fields 481 | 482 | def _set_fields(self, val): 483 | self._validate_option("fields", val) 484 | self._fields = val 485 | 486 | fields = property(_get_fields, _set_fields) 487 | 488 | def _get_start(self): 489 | """Start index of the range of rows to print 490 | 491 | Arguments: 492 | 493 | start - index of first data row to include in output""" 494 | return self._start 495 | 496 | def _set_start(self, val): 497 | self._validate_option("start", val) 498 | self._start = val 499 | 500 | start = property(_get_start, _set_start) 501 | 502 | def _get_end(self): 503 | """End index of the range of rows to print 504 | 505 | Arguments: 506 | 507 | end - index of last data row to include in output PLUS ONE (list slice style)""" 508 | return self._end 509 | 510 | def _set_end(self, val): 511 | self._validate_option("end", val) 512 | self._end = val 513 | 514 | end = property(_get_end, _set_end) 515 | 516 | def _get_sortby(self): 517 | """Name of field by which to sort rows 518 | 519 | Arguments: 520 | 521 | sortby - field name to sort by""" 522 | return self._sortby 523 | 524 | def _set_sortby(self, val): 525 | self._validate_option("sortby", val) 526 | self._sortby = val 527 | 528 | sortby = property(_get_sortby, _set_sortby) 529 | 530 | def _get_reversesort(self): 531 | """Controls direction of sorting (ascending vs descending) 532 | 533 | Arguments: 534 | 535 | reveresort - set to True to sort by descending order, or False to sort by ascending order""" 536 | return self._reversesort 537 | 538 | def _set_reversesort(self, val): 539 | self._validate_option("reversesort", val) 540 | self._reversesort = val 541 | 542 | reversesort = property(_get_reversesort, _set_reversesort) 543 | 544 | def _get_sort_key(self): 545 | """Sorting key function, applied to data points before sorting 546 | 547 | Arguments: 548 | 549 | sort_key - a function which takes one argument and returns something to be sorted""" 550 | return self._sort_key 551 | 552 | def _set_sort_key(self, val): 553 | self._validate_option("sort_key", val) 554 | self._sort_key = val 555 | 556 | sort_key = property(_get_sort_key, _set_sort_key) 557 | 558 | def _get_header(self): 559 | """Controls printing of table header with field names 560 | 561 | Arguments: 562 | 563 | header - print a header showing field names (True or False)""" 564 | return self._header 565 | 566 | def _set_header(self, val): 567 | self._validate_option("header", val) 568 | self._header = val 569 | 570 | header = property(_get_header, _set_header) 571 | 572 | def _get_header_style(self): 573 | """Controls stylisation applied to field names in header 574 | 575 | Arguments: 576 | 577 | header_style - stylisation to apply to field names in header ("cap", "title", "upper", "lower" or None)""" 578 | return self._header_style 579 | 580 | def _set_header_style(self, val): 581 | self._validate_header_style(val) 582 | self._header_style = val 583 | 584 | header_style = property(_get_header_style, _set_header_style) 585 | 586 | def _get_border(self): 587 | """Controls printing of border around table 588 | 589 | Arguments: 590 | 591 | border - print a border around the table (True or False)""" 592 | return self._border 593 | 594 | def _set_border(self, val): 595 | self._validate_option("border", val) 596 | self._border = val 597 | 598 | border = property(_get_border, _set_border) 599 | 600 | def _get_hrules(self): 601 | """Controls printing of horizontal rules after rows 602 | 603 | Arguments: 604 | 605 | hrules - horizontal rules style. Allowed values: FRAME, ALL, HEADER, NONE""" 606 | return self._hrules 607 | 608 | def _set_hrules(self, val): 609 | self._validate_option("hrules", val) 610 | self._hrules = val 611 | 612 | hrules = property(_get_hrules, _set_hrules) 613 | 614 | def _get_vrules(self): 615 | """Controls printing of vertical rules between columns 616 | 617 | Arguments: 618 | 619 | vrules - vertical rules style. Allowed values: FRAME, ALL, NONE""" 620 | return self._vrules 621 | 622 | def _set_vrules(self, val): 623 | self._validate_option("vrules", val) 624 | self._vrules = val 625 | 626 | vrules = property(_get_vrules, _set_vrules) 627 | 628 | def _get_int_format(self): 629 | """Controls formatting of integer data 630 | Arguments: 631 | 632 | int_format - integer format string""" 633 | return self._int_format 634 | 635 | def _set_int_format(self, val): 636 | # self._validate_option("int_format", val) 637 | for field in self._field_names: 638 | self._int_format[field] = val 639 | 640 | int_format = property(_get_int_format, _set_int_format) 641 | 642 | def _get_float_format(self): 643 | """Controls formatting of floating point data 644 | Arguments: 645 | 646 | float_format - floating point format string""" 647 | return self._float_format 648 | 649 | def _set_float_format(self, val): 650 | # self._validate_option("float_format", val) 651 | for field in self._field_names: 652 | self._float_format[field] = val 653 | 654 | float_format = property(_get_float_format, _set_float_format) 655 | 656 | def _get_padding_width(self): 657 | """The number of empty spaces between a column's edge and its content 658 | 659 | Arguments: 660 | 661 | padding_width - number of spaces, must be a positive integer""" 662 | return self._padding_width 663 | 664 | def _set_padding_width(self, val): 665 | self._validate_option("padding_width", val) 666 | self._padding_width = val 667 | 668 | padding_width = property(_get_padding_width, _set_padding_width) 669 | 670 | def _get_left_padding_width(self): 671 | """The number of empty spaces between a column's left edge and its content 672 | 673 | Arguments: 674 | 675 | left_padding - number of spaces, must be a positive integer""" 676 | return self._left_padding_width 677 | 678 | def _set_left_padding_width(self, val): 679 | self._validate_option("left_padding_width", val) 680 | self._left_padding_width = val 681 | 682 | left_padding_width = property(_get_left_padding_width, _set_left_padding_width) 683 | 684 | def _get_right_padding_width(self): 685 | """The number of empty spaces between a column's right edge and its content 686 | 687 | Arguments: 688 | 689 | right_padding - number of spaces, must be a positive integer""" 690 | return self._right_padding_width 691 | 692 | def _set_right_padding_width(self, val): 693 | self._validate_option("right_padding_width", val) 694 | self._right_padding_width = val 695 | 696 | right_padding_width = property(_get_right_padding_width, _set_right_padding_width) 697 | 698 | def _get_vertical_char(self): 699 | """The charcter used when printing table borders to draw vertical lines 700 | 701 | Arguments: 702 | 703 | vertical_char - single character string used to draw vertical lines""" 704 | return self._vertical_char 705 | 706 | def _set_vertical_char(self, val): 707 | val = self._unicode(val) 708 | self._validate_option("vertical_char", val) 709 | self._vertical_char = val 710 | 711 | vertical_char = property(_get_vertical_char, _set_vertical_char) 712 | 713 | def _get_horizontal_char(self): 714 | """The charcter used when printing table borders to draw horizontal lines 715 | 716 | Arguments: 717 | 718 | horizontal_char - single character string used to draw horizontal lines""" 719 | return self._horizontal_char 720 | 721 | def _set_horizontal_char(self, val): 722 | val = self._unicode(val) 723 | self._validate_option("horizontal_char", val) 724 | self._horizontal_char = val 725 | 726 | horizontal_char = property(_get_horizontal_char, _set_horizontal_char) 727 | 728 | def _get_junction_char(self): 729 | """The charcter used when printing table borders to draw line junctions 730 | 731 | Arguments: 732 | 733 | junction_char - single character string used to draw line junctions""" 734 | return self._junction_char 735 | 736 | def _set_junction_char(self, val): 737 | val = self._unicode(val) 738 | self._validate_option("vertical_char", val) 739 | self._junction_char = val 740 | 741 | junction_char = property(_get_junction_char, _set_junction_char) 742 | 743 | def _get_format(self): 744 | """Controls whether or not HTML tables are formatted to match styling options 745 | 746 | Arguments: 747 | 748 | format - True or False""" 749 | return self._format 750 | 751 | def _set_format(self, val): 752 | self._validate_option("format", val) 753 | self._format = val 754 | 755 | format = property(_get_format, _set_format) 756 | 757 | def _get_print_empty(self): 758 | """Controls whether or not empty tables produce a header and frame or just an empty string 759 | 760 | Arguments: 761 | 762 | print_empty - True or False""" 763 | return self._print_empty 764 | 765 | def _set_print_empty(self, val): 766 | self._validate_option("print_empty", val) 767 | self._print_empty = val 768 | 769 | print_empty = property(_get_print_empty, _set_print_empty) 770 | 771 | def _get_attributes(self): 772 | """A dictionary of HTML attribute name/value pairs to be included in the tag when printing HTML 773 | 774 | Arguments: 775 | 776 | attributes - dictionary of attributes""" 777 | return self._attributes 778 | 779 | def _set_attributes(self, val): 780 | self._validate_option("attributes", val) 781 | self._attributes = val 782 | 783 | attributes = property(_get_attributes, _set_attributes) 784 | 785 | ############################## 786 | # OPTION MIXER # 787 | ############################## 788 | 789 | def _get_options(self, kwargs): 790 | 791 | options = {} 792 | for option in self._options: 793 | if option in kwargs: 794 | self._validate_option(option, kwargs[option]) 795 | options[option] = kwargs[option] 796 | else: 797 | options[option] = getattr(self, "_" + option) 798 | return options 799 | 800 | ############################## 801 | # PRESET STYLE LOGIC # 802 | ############################## 803 | 804 | def set_style(self, style): 805 | 806 | if style == DEFAULT: 807 | self._set_default_style() 808 | elif style == MSWORD_FRIENDLY: 809 | self._set_msword_style() 810 | elif style == PLAIN_COLUMNS: 811 | self._set_columns_style() 812 | elif style == RANDOM: 813 | self._set_random_style() 814 | else: 815 | raise Exception("Invalid pre-set style!") 816 | 817 | def _set_default_style(self): 818 | 819 | self.header = True 820 | self.border = True 821 | self._hrules = FRAME 822 | self._vrules = ALL 823 | self.padding_width = 1 824 | self.left_padding_width = 1 825 | self.right_padding_width = 1 826 | self.vertical_char = "|" 827 | self.horizontal_char = "-" 828 | self.junction_char = "+" 829 | 830 | def _set_msword_style(self): 831 | 832 | self.header = True 833 | self.border = True 834 | self._hrules = NONE 835 | self.padding_width = 1 836 | self.left_padding_width = 1 837 | self.right_padding_width = 1 838 | self.vertical_char = "|" 839 | 840 | def _set_columns_style(self): 841 | 842 | self.header = True 843 | self.border = False 844 | self.padding_width = 1 845 | self.left_padding_width = 0 846 | self.right_padding_width = 8 847 | 848 | def _set_random_style(self): 849 | 850 | # Just for fun! 851 | self.header = random.choice((True, False)) 852 | self.border = random.choice((True, False)) 853 | self._hrules = random.choice((ALL, FRAME, HEADER, NONE)) 854 | self._vrules = random.choice((ALL, FRAME, NONE)) 855 | self.left_padding_width = random.randint(0, 5) 856 | self.right_padding_width = random.randint(0, 5) 857 | self.vertical_char = random.choice("~!@#$%^&*()_+|-=\{}[];':\",./;<>?") 858 | self.horizontal_char = random.choice("~!@#$%^&*()_+|-=\{}[];':\",./;<>?") 859 | self.junction_char = random.choice("~!@#$%^&*()_+|-=\{}[];':\",./;<>?") 860 | 861 | ############################## 862 | # DATA INPUT METHODS # 863 | ############################## 864 | 865 | def add_row(self, row): 866 | 867 | """Add a row to the table 868 | 869 | Arguments: 870 | 871 | row - row of data, should be a list with as many elements as the table 872 | has fields""" 873 | 874 | if self._field_names and len(row) != len(self._field_names): 875 | raise Exception( 876 | "Row has incorrect number of values, (actual) %d!=%d (expected)" % (len(row), len(self._field_names))) 877 | if not self._field_names: 878 | self.field_names = [("Field %d" % (n + 1)) for n in range(0, len(row))] 879 | self._rows.append(list(row)) 880 | 881 | def del_row(self, row_index): 882 | 883 | """Delete a row to the table 884 | 885 | Arguments: 886 | 887 | row_index - The index of the row you want to delete. Indexing starts at 0.""" 888 | 889 | if row_index > len(self._rows) - 1: 890 | raise Exception("Cant delete row at index %d, table only has %d rows!" % (row_index, len(self._rows))) 891 | del self._rows[row_index] 892 | 893 | def add_column(self, fieldname, column, align="c", valign="t"): 894 | 895 | """Add a column to the table. 896 | 897 | Arguments: 898 | 899 | fieldname - name of the field to contain the new column of data 900 | column - column of data, should be a list with as many elements as the 901 | table has rows 902 | align - desired alignment for this column - "l" for left, "c" for centre and "r" for right 903 | valign - desired vertical alignment for new columns - "t" for top, "m" for middle and "b" for bottom""" 904 | 905 | if len(self._rows) in (0, len(column)): 906 | self._validate_align(align) 907 | self._validate_valign(valign) 908 | self._field_names.append(fieldname) 909 | self._align[fieldname] = align 910 | self._valign[fieldname] = valign 911 | for i in range(0, len(column)): 912 | if len(self._rows) < i + 1: 913 | self._rows.append([]) 914 | self._rows[i].append(column[i]) 915 | else: 916 | raise Exception("Column length %d does not match number of rows %d!" % (len(column), len(self._rows))) 917 | 918 | def clear_rows(self): 919 | 920 | """Delete all rows from the table but keep the current field names""" 921 | 922 | self._rows = [] 923 | 924 | def clear(self): 925 | 926 | """Delete all rows and field names from the table, maintaining nothing but styling options""" 927 | 928 | self._rows = [] 929 | self._field_names = [] 930 | self._widths = [] 931 | 932 | ############################## 933 | # MISC PUBLIC METHODS # 934 | ############################## 935 | 936 | def copy(self): 937 | return copy.deepcopy(self) 938 | 939 | ############################## 940 | # MISC PRIVATE METHODS # 941 | ############################## 942 | 943 | def _format_value(self, field, value): 944 | if isinstance(value, int) and field in self._int_format: 945 | value = self._unicode(("%%%sd" % self._int_format[field]) % value) 946 | elif isinstance(value, float) and field in self._float_format: 947 | value = self._unicode(("%%%sf" % self._float_format[field]) % value) 948 | return self._unicode(value) 949 | 950 | def _compute_widths(self, rows, options): 951 | if options["header"]: 952 | widths = [_get_size(field)[0] for field in self._field_names] 953 | else: 954 | widths = len(self.field_names) * [0] 955 | for row in rows: 956 | for index, value in enumerate(row): 957 | fieldname = self.field_names[index] 958 | if fieldname in self.max_width: 959 | widths[index] = max(widths[index], min(_get_size(value)[0], self.max_width[fieldname])) 960 | else: 961 | widths[index] = max(widths[index], _get_size(value)[0]) 962 | self._widths = widths 963 | 964 | def _get_padding_widths(self, options): 965 | 966 | if options["left_padding_width"] is not None: 967 | lpad = options["left_padding_width"] 968 | else: 969 | lpad = options["padding_width"] 970 | if options["right_padding_width"] is not None: 971 | rpad = options["right_padding_width"] 972 | else: 973 | rpad = options["padding_width"] 974 | return lpad, rpad 975 | 976 | def _get_rows(self, options): 977 | """Return only those data rows that should be printed, based on slicing and sorting. 978 | 979 | Arguments: 980 | 981 | options - dictionary of option settings.""" 982 | 983 | # Make a copy of only those rows in the slice range 984 | rows = copy.deepcopy(self._rows[options["start"]:options["end"]]) 985 | # Sort if necessary 986 | if options["sortby"]: 987 | sortindex = self._field_names.index(options["sortby"]) 988 | # Decorate 989 | rows = [[row[sortindex]] + row for row in rows] 990 | # Sort 991 | rows.sort(reverse=options["reversesort"], key=options["sort_key"]) 992 | # Undecorate 993 | rows = [row[1:] for row in rows] 994 | return rows 995 | 996 | def _format_row(self, row, options): 997 | return [self._format_value(field, value) for (field, value) in zip(self._field_names, row)] 998 | 999 | def _format_rows(self, rows, options): 1000 | return [self._format_row(row, options) for row in rows] 1001 | 1002 | ############################## 1003 | # PLAIN TEXT STRING METHODS # 1004 | ############################## 1005 | 1006 | def get_string(self, **kwargs): 1007 | 1008 | """Return string representation of table in current state. 1009 | 1010 | Arguments: 1011 | 1012 | start - index of first data row to include in output 1013 | end - index of last data row to include in output PLUS ONE (list slice style) 1014 | fields - names of fields (columns) to include 1015 | header - print a header showing field names (True or False) 1016 | border - print a border around the table (True or False) 1017 | hrules - controls printing of horizontal rules after rows. Allowed values: ALL, FRAME, HEADER, NONE 1018 | vrules - controls printing of vertical rules between columns. Allowed values: FRAME, ALL, NONE 1019 | int_format - controls formatting of integer data 1020 | float_format - controls formatting of floating point data 1021 | padding_width - number of spaces on either side of column data (only used if left and right paddings are None) 1022 | left_padding_width - number of spaces on left hand side of column data 1023 | right_padding_width - number of spaces on right hand side of column data 1024 | vertical_char - single character string used to draw vertical lines 1025 | horizontal_char - single character string used to draw horizontal lines 1026 | junction_char - single character string used to draw line junctions 1027 | sortby - name of field to sort rows by 1028 | sort_key - sorting key function, applied to data points before sorting 1029 | reversesort - True or False to sort in descending or ascending order 1030 | print empty - if True, stringify just the header for an empty table, if False return an empty string """ 1031 | 1032 | options = self._get_options(kwargs) 1033 | 1034 | lines = [] 1035 | 1036 | # Don't think too hard about an empty table 1037 | # Is this the desired behaviour? Maybe we should still print the header? 1038 | if self.rowcount == 0 and (not options["print_empty"] or not options["border"]): 1039 | return "" 1040 | 1041 | # Get the rows we need to print, taking into account slicing, sorting, etc. 1042 | rows = self._get_rows(options) 1043 | 1044 | # Turn all data in all rows into Unicode, formatted as desired 1045 | formatted_rows = self._format_rows(rows, options) 1046 | 1047 | # Compute column widths 1048 | self._compute_widths(formatted_rows, options) 1049 | 1050 | # Add header or top of border 1051 | self._hrule = self._stringify_hrule(options) 1052 | if options["header"]: 1053 | lines.append(self._stringify_header(options)) 1054 | elif options["border"] and options["hrules"] in (ALL, FRAME): 1055 | lines.append(self._hrule) 1056 | 1057 | # Add rows 1058 | for row in formatted_rows: 1059 | lines.append(self._stringify_row(row, options)) 1060 | 1061 | # Add bottom of border 1062 | if options["border"] and options["hrules"] == FRAME: 1063 | lines.append(self._hrule) 1064 | 1065 | return self._unicode("\n").join(lines) 1066 | 1067 | def _stringify_hrule(self, options): 1068 | 1069 | if not options["border"]: 1070 | return "" 1071 | lpad, rpad = self._get_padding_widths(options) 1072 | if options['vrules'] in (ALL, FRAME): 1073 | bits = [options["junction_char"]] 1074 | else: 1075 | bits = [options["horizontal_char"]] 1076 | # For tables with no data or fieldnames 1077 | if not self._field_names: 1078 | bits.append(options["junction_char"]) 1079 | return "".join(bits) 1080 | for field, width in zip(self._field_names, self._widths): 1081 | if options["fields"] and field not in options["fields"]: 1082 | continue 1083 | bits.append((width + lpad + rpad) * options["horizontal_char"]) 1084 | if options['vrules'] == ALL: 1085 | bits.append(options["junction_char"]) 1086 | else: 1087 | bits.append(options["horizontal_char"]) 1088 | if options["vrules"] == FRAME: 1089 | bits.pop() 1090 | bits.append(options["junction_char"]) 1091 | return "".join(bits) 1092 | 1093 | def _stringify_header(self, options): 1094 | 1095 | bits = [] 1096 | lpad, rpad = self._get_padding_widths(options) 1097 | if options["border"]: 1098 | if options["hrules"] in (ALL, FRAME): 1099 | bits.append(self._hrule) 1100 | bits.append("\n") 1101 | if options["vrules"] in (ALL, FRAME): 1102 | bits.append(options["vertical_char"]) 1103 | else: 1104 | bits.append(" ") 1105 | # For tables with no data or field names 1106 | if not self._field_names: 1107 | if options["vrules"] in (ALL, FRAME): 1108 | bits.append(options["vertical_char"]) 1109 | else: 1110 | bits.append(" ") 1111 | for field, width, in zip(self._field_names, self._widths): 1112 | if options["fields"] and field not in options["fields"]: 1113 | continue 1114 | if self._header_style == "cap": 1115 | fieldname = field.capitalize() 1116 | elif self._header_style == "title": 1117 | fieldname = field.title() 1118 | elif self._header_style == "upper": 1119 | fieldname = field.upper() 1120 | elif self._header_style == "lower": 1121 | fieldname = field.lower() 1122 | else: 1123 | fieldname = field 1124 | bits.append(" " * lpad + self._justify(fieldname, width, self._align[field]) + " " * rpad) 1125 | if options["border"]: 1126 | if options["vrules"] == ALL: 1127 | bits.append(options["vertical_char"]) 1128 | else: 1129 | bits.append(" ") 1130 | # If vrules is FRAME, then we just appended a space at the end 1131 | # of the last field, when we really want a vertical character 1132 | if options["border"] and options["vrules"] == FRAME: 1133 | bits.pop() 1134 | bits.append(options["vertical_char"]) 1135 | if options["border"] and options["hrules"] != NONE: 1136 | bits.append("\n") 1137 | bits.append(self._hrule) 1138 | return "".join(bits) 1139 | 1140 | def _stringify_row(self, row, options): 1141 | 1142 | for index, field, value, width, in zip(range(0, len(row)), self._field_names, row, self._widths): 1143 | # Enforce max widths 1144 | lines = value.split("\n") 1145 | new_lines = [] 1146 | for line in lines: 1147 | if _str_block_width(line) > width: 1148 | line = textwrap.fill(line, width) 1149 | new_lines.append(line) 1150 | lines = new_lines 1151 | value = "\n".join(lines) 1152 | row[index] = value 1153 | 1154 | row_height = 0 1155 | for c in row: 1156 | h = _get_size(c)[1] 1157 | if h > row_height: 1158 | row_height = h 1159 | 1160 | bits = [] 1161 | lpad, rpad = self._get_padding_widths(options) 1162 | for y in range(0, row_height): 1163 | bits.append([]) 1164 | if options["border"]: 1165 | if options["vrules"] in (ALL, FRAME): 1166 | bits[y].append(self.vertical_char) 1167 | else: 1168 | bits[y].append(" ") 1169 | 1170 | for field, value, width, in zip(self._field_names, row, self._widths): 1171 | 1172 | valign = self._valign[field] 1173 | lines = value.split("\n") 1174 | dHeight = row_height - len(lines) 1175 | if dHeight: 1176 | if valign == "m": 1177 | lines = [""] * int(dHeight / 2) + lines + [""] * (dHeight - int(dHeight / 2)) 1178 | elif valign == "b": 1179 | lines = [""] * dHeight + lines 1180 | else: 1181 | lines = lines + [""] * dHeight 1182 | 1183 | y = 0 1184 | for l in lines: 1185 | if options["fields"] and field not in options["fields"]: 1186 | continue 1187 | 1188 | bits[y].append(" " * lpad + self._justify(l, width, self._align[field]) + " " * rpad) 1189 | if options["border"]: 1190 | if options["vrules"] == ALL: 1191 | bits[y].append(self.vertical_char) 1192 | else: 1193 | bits[y].append(" ") 1194 | y += 1 1195 | 1196 | # If vrules is FRAME, then we just appended a space at the end 1197 | # of the last field, when we really want a vertical character 1198 | for y in range(0, row_height): 1199 | if options["border"] and options["vrules"] == FRAME: 1200 | bits[y].pop() 1201 | bits[y].append(options["vertical_char"]) 1202 | 1203 | if options["border"] and options["hrules"] == ALL: 1204 | bits[row_height - 1].append("\n") 1205 | bits[row_height - 1].append(self._hrule) 1206 | 1207 | for y in range(0, row_height): 1208 | bits[y] = "".join(bits[y]) 1209 | 1210 | return "\n".join(bits) 1211 | 1212 | ############################## 1213 | # HTML STRING METHODS # 1214 | ############################## 1215 | 1216 | def get_html_string(self, **kwargs): 1217 | 1218 | """Return string representation of HTML formatted version of table in current state. 1219 | 1220 | Arguments: 1221 | 1222 | start - index of first data row to include in output 1223 | end - index of last data row to include in output PLUS ONE (list slice style) 1224 | fields - names of fields (columns) to include 1225 | header - print a header showing field names (True or False) 1226 | border - print a border around the table (True or False) 1227 | hrules - controls printing of horizontal rules after rows. Allowed values: ALL, FRAME, HEADER, NONE 1228 | vrules - controls printing of vertical rules between columns. Allowed values: FRAME, ALL, NONE 1229 | int_format - controls formatting of integer data 1230 | float_format - controls formatting of floating point data 1231 | padding_width - number of spaces on either side of column data (only used if left and right paddings are None) 1232 | left_padding_width - number of spaces on left hand side of column data 1233 | right_padding_width - number of spaces on right hand side of column data 1234 | sortby - name of field to sort rows by 1235 | sort_key - sorting key function, applied to data points before sorting 1236 | attributes - dictionary of name/value pairs to include as HTML attributes in the
tag 1237 | xhtml - print
tags if True,
tags if false""" 1238 | 1239 | options = self._get_options(kwargs) 1240 | 1241 | if options["format"]: 1242 | string = self._get_formatted_html_string(options) 1243 | else: 1244 | string = self._get_simple_html_string(options) 1245 | 1246 | return string 1247 | 1248 | def _get_simple_html_string(self, options): 1249 | 1250 | lines = [] 1251 | if options["xhtml"]: 1252 | linebreak = "
" 1253 | else: 1254 | linebreak = "
" 1255 | 1256 | open_tag = [] 1257 | open_tag.append("") 1262 | lines.append("".join(open_tag)) 1263 | 1264 | # Headers 1265 | if options["header"]: 1266 | lines.append(" ") 1267 | for field in self._field_names: 1268 | if options["fields"] and field not in options["fields"]: 1269 | continue 1270 | lines.append(" " % escape(field).replace("\n", linebreak)) 1271 | lines.append(" ") 1272 | 1273 | # Data 1274 | rows = self._get_rows(options) 1275 | formatted_rows = self._format_rows(rows, options) 1276 | for row in formatted_rows: 1277 | lines.append(" ") 1278 | for field, datum in zip(self._field_names, row): 1279 | if options["fields"] and field not in options["fields"]: 1280 | continue 1281 | lines.append(" " % escape(datum).replace("\n", linebreak)) 1282 | lines.append(" ") 1283 | 1284 | lines.append("
%s
%s
") 1285 | 1286 | return self._unicode("\n").join(lines) 1287 | 1288 | def _get_formatted_html_string(self, options): 1289 | 1290 | lines = [] 1291 | lpad, rpad = self._get_padding_widths(options) 1292 | if options["xhtml"]: 1293 | linebreak = "
" 1294 | else: 1295 | linebreak = "
" 1296 | 1297 | open_tag = [] 1298 | open_tag.append("") 1318 | lines.append("".join(open_tag)) 1319 | 1320 | # Headers 1321 | if options["header"]: 1322 | lines.append(" ") 1323 | for field in self._field_names: 1324 | if options["fields"] and field not in options["fields"]: 1325 | continue 1326 | lines.append( 1327 | " %s" % ( 1328 | lpad, rpad, escape(field).replace("\n", linebreak))) 1329 | lines.append(" ") 1330 | 1331 | # Data 1332 | rows = self._get_rows(options) 1333 | formatted_rows = self._format_rows(rows, options) 1334 | aligns = [] 1335 | valigns = [] 1336 | for field in self._field_names: 1337 | aligns.append({"l": "left", "r": "right", "c": "center"}[self._align[field]]) 1338 | valigns.append({"t": "top", "m": "middle", "b": "bottom"}[self._valign[field]]) 1339 | for row in formatted_rows: 1340 | lines.append(" ") 1341 | for field, datum, align, valign in zip(self._field_names, row, aligns, valigns): 1342 | if options["fields"] and field not in options["fields"]: 1343 | continue 1344 | lines.append( 1345 | " %s" % ( 1346 | lpad, rpad, align, valign, escape(datum).replace("\n", linebreak))) 1347 | lines.append(" ") 1348 | lines.append("") 1349 | 1350 | return self._unicode("\n").join(lines) 1351 | 1352 | 1353 | ############################## 1354 | # UNICODE WIDTH FUNCTIONS # 1355 | ############################## 1356 | 1357 | def _char_block_width(char): 1358 | # Basic Latin, which is probably the most common case 1359 | # if char in xrange(0x0021, 0x007e): 1360 | # if char >= 0x0021 and char <= 0x007e: 1361 | if 0x0021 <= char <= 0x007e: 1362 | return 1 1363 | # Chinese, Japanese, Korean (common) 1364 | if 0x4e00 <= char <= 0x9fff: 1365 | return 2 1366 | # Hangul 1367 | if 0xac00 <= char <= 0xd7af: 1368 | return 2 1369 | # Combining? 1370 | if unicodedata.combining(uni_chr(char)): 1371 | return 0 1372 | # Hiragana and Katakana 1373 | if 0x3040 <= char <= 0x309f or 0x30a0 <= char <= 0x30ff: 1374 | return 2 1375 | # Full-width Latin characters 1376 | if 0xff01 <= char <= 0xff60: 1377 | return 2 1378 | # CJK punctuation 1379 | if 0x3000 <= char <= 0x303e: 1380 | return 2 1381 | # Backspace and delete 1382 | if char in (0x0008, 0x007f): 1383 | return -1 1384 | # Other control characters 1385 | elif char in (0x0000, 0x001f): 1386 | return 0 1387 | # Take a guess 1388 | return 1 1389 | 1390 | 1391 | def _str_block_width(val): 1392 | return sum(itermap(_char_block_width, itermap(ord, _re.sub("", val)))) 1393 | 1394 | 1395 | ############################## 1396 | # TABLE FACTORIES # 1397 | ############################## 1398 | 1399 | def from_csv(fp, field_names=None, **kwargs): 1400 | dialect = csv.Sniffer().sniff(fp.read(1024)) 1401 | fp.seek(0) 1402 | reader = csv.reader(fp, dialect) 1403 | 1404 | table = PrettyTable(**kwargs) 1405 | if field_names: 1406 | table.field_names = field_names 1407 | else: 1408 | if py3k: 1409 | table.field_names = [x.strip() for x in next(reader)] 1410 | else: 1411 | table.field_names = [x.strip() for x in reader.next()] 1412 | 1413 | for row in reader: 1414 | table.add_row([x.strip() for x in row]) 1415 | 1416 | return table 1417 | 1418 | 1419 | def from_db_cursor(cursor, **kwargs): 1420 | if cursor.description: 1421 | table = PrettyTable(**kwargs) 1422 | table.field_names = [col[0] for col in cursor.description] 1423 | for row in cursor.fetchall(): 1424 | table.add_row(row) 1425 | return table 1426 | 1427 | 1428 | class TableHandler(HTMLParser): 1429 | 1430 | def __init__(self, **kwargs): 1431 | HTMLParser.__init__(self) 1432 | self.kwargs = kwargs 1433 | self.tables = [] 1434 | self.last_row = [] 1435 | self.rows = [] 1436 | self.max_row_width = 0 1437 | self.active = None 1438 | self.last_content = "" 1439 | self.is_last_row_header = False 1440 | 1441 | def handle_starttag(self, tag, attrs): 1442 | self.active = tag 1443 | if tag == "th": 1444 | self.is_last_row_header = True 1445 | 1446 | def handle_endtag(self, tag): 1447 | if tag in ["th", "td"]: 1448 | stripped_content = self.last_content.strip() 1449 | self.last_row.append(stripped_content) 1450 | if tag == "tr": 1451 | self.rows.append( 1452 | (self.last_row, self.is_last_row_header)) 1453 | self.max_row_width = max(self.max_row_width, len(self.last_row)) 1454 | self.last_row = [] 1455 | self.is_last_row_header = False 1456 | if tag == "table": 1457 | table = self.generate_table(self.rows) 1458 | self.tables.append(table) 1459 | self.rows = [] 1460 | self.last_content = " " 1461 | self.active = None 1462 | 1463 | def handle_data(self, data): 1464 | self.last_content += data 1465 | 1466 | def generate_table(self, rows): 1467 | """ 1468 | Generates from a list of rows a PrettyTable object. 1469 | """ 1470 | table = PrettyTable(**self.kwargs) 1471 | for row in self.rows: 1472 | if len(row[0]) < self.max_row_width: 1473 | appends = self.max_row_width - len(row[0]) 1474 | for i in range(1, appends): 1475 | row[0].append("-") 1476 | 1477 | if row[1] == True: 1478 | self.make_fields_unique(row[0]) 1479 | table.field_names = row[0] 1480 | else: 1481 | table.add_row(row[0]) 1482 | return table 1483 | 1484 | def make_fields_unique(self, fields): 1485 | """ 1486 | iterates over the row and make each field unique 1487 | """ 1488 | for i in range(0, len(fields)): 1489 | for j in range(i + 1, len(fields)): 1490 | if fields[i] == fields[j]: 1491 | fields[j] += "'" 1492 | 1493 | 1494 | def from_html(html_code, **kwargs): 1495 | """ 1496 | Generates a list of PrettyTables from a string of HTML code. Each in 1497 | the HTML becomes one PrettyTable object. 1498 | """ 1499 | 1500 | parser = TableHandler(**kwargs) 1501 | parser.feed(html_code) 1502 | return parser.tables 1503 | 1504 | 1505 | def from_html_one(html_code, **kwargs): 1506 | """ 1507 | Generates a PrettyTables from a string of HTML code which contains only a 1508 | single
1509 | """ 1510 | 1511 | tables = from_html(html_code, **kwargs) 1512 | try: 1513 | assert len(tables) == 1 1514 | except AssertionError: 1515 | raise Exception("More than one
in provided HTML code! Use from_html instead.") 1516 | return tables[0] 1517 | 1518 | 1519 | ############################## 1520 | # MAIN (TEST FUNCTION) # 1521 | ############################## 1522 | 1523 | def main(): 1524 | x = PrettyTable(["City name", "Area", "Population", "Annual Rainfall"]) 1525 | x.sortby = "Population" 1526 | x.reversesort = True 1527 | x.int_format["Area"] = "04d" 1528 | x.float_format = "6.1f" 1529 | x.align["City name"] = "l" # Left align city names 1530 | x.add_row(["Adelaide", 1295, 1158259, 600.5]) 1531 | x.add_row(["Brisbane", 5905, 1857594, 1146.4]) 1532 | x.add_row(["Darwin", 112, 120900, 1714.7]) 1533 | x.add_row(["Hobart", 1357, 205556, 619.5]) 1534 | x.add_row(["Sydney", 2058, 4336374, 1214.8]) 1535 | x.add_row(["Melbourne", 1566, 3806092, 646.9]) 1536 | x.add_row(["Perth", 5386, 1554769, 869.4]) 1537 | print(x) 1538 | 1539 | 1540 | if __name__ == "__main__": 1541 | main() --------------------------------------------------------------------------------