-
Notifications
You must be signed in to change notification settings - Fork 523
/
Copy paththrift_printer.py
153 lines (137 loc) · 6.25 KB
/
thrift_printer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
#
import re
from sys import stdout
class ThriftPrettyPrinter(object):
"""Implements a pretty printer for Thrift objects.
Generates string representations. Does not output
to stdout, stderr, or any other file handles."""
# Inputs:
# redacted_fields - list of names of object attributes whose
# values will not be printed out
# indent - string containing only spaces, used as the
# base indentation where all other indentations
# will be multiples of this string
# objects_to_skip - list of names of objects attributes that
# will not be printed out, useful to ignore
# duplicate information such as query results
def __init__(self,
redacted_fields=("secret", "password"),
indent=" ",
objects_to_skip=("TRowSet", "TGetRuntimeProfileResp")):
if redacted_fields is not None:
assert type(redacted_fields) is list or \
type(redacted_fields) is tuple, \
"redacted_fields must be either a list or a tuple"
self.base_indent = indent
self.redacted_fields = redacted_fields
self.objects_to_skip = objects_to_skip
self._objname_re = re.compile("^.*?'(.*?)'.*?$")
def print_obj(self, thrift_obj, file_handle=stdout):
"""Prints the provided 'thrift_obj' to the provided
file handle. If no file handle is specified, then
stdout will be used. The 'file_handle' object must
have a write(string) method.
While this class is specifically targeted to printing
Thrift objects, there is no technical limitation preventing
any other type of object from being printed. However, the
output of non-Thrift objects may not be as nicely formatted."""
# Inputs:
# thrift_obj - the object to print out, its attributes will
# be walked recursively through the entire
# object structure and printed out
# file_handle - where the object will be written, defaults to stdout
# but can be any object with a write(str) method
self._internal_print(thrift_obj, self.base_indent, file_handle)
def _internal_print(self, thrift_obj, indent, file_handle):
"""Recursive function that does the work of walking and printing
an object."""
# parse out the type name of the thrift object
obj_name = self._objname_re.match(str(type(thrift_obj))) \
.group(1).split(".")[-1]
file_handle.write("<{0}>".format(obj_name))
if self.objects_to_skip.count(obj_name):
file_handle.write(" - <skipping>\n")
return
indent = "{0}{1}".format(indent, self.base_indent)
file_handle.write("\n")
if obj_name == "list" or obj_name == "tuple":
# lists and tuples have to be handled differently
# because the vars function does not operate on them
for attr_val in thrift_obj:
file_handle.write(indent)
self._internal_print(attr_val, indent, file_handle)
else:
# print out simple types first before printing out objects
# this ensures the simple types are easier to see
child_simple_attrs = {}
child_objs = {}
for attr_name in vars(thrift_obj):
attr_val = getattr(thrift_obj, attr_name)
if (hasattr(attr_val, '__dict__')
or attr_val is list
or attr_val is tuple):
child_objs[attr_name] = attr_val
else:
child_simple_attrs[attr_name] = attr_val
# print out child attributes in alphabetical order
for child_attr_name in sorted(child_simple_attrs):
self._print_attr(child_attr_name,
child_simple_attrs[child_attr_name],
indent,
file_handle)
# print out complex types objects, lists, or tuples
# in alphabetical order
for attr_name in sorted(child_objs):
self._print_attr(attr_name,
child_objs[attr_name],
indent,
file_handle)
def _print_attr(self, attr_name, attr_val, indent, file_handle):
"""Handles a single object attribute by either printing out
its name/value (for simple types) or recursing down into the
object (for objects/lists/tuples)."""
file_handle.write(indent)
if attr_val is not None and self.redacted_fields.count(attr_name) > 0:
file_handle.write("- {0}: *******\n".format(attr_name))
elif attr_val is None:
file_handle.write("- {0}: <None>\n".format(attr_name))
elif type(attr_val) is list or type(attr_val) is tuple:
file_handle.write("[")
self._internal_print(attr_val, indent, file_handle)
file_handle.write("{0}]\n".format(indent))
elif hasattr(attr_val, '__dict__'):
indent += "{0:{1}} {2}".format("", len(attr_name), self.base_indent)
file_handle.write("- {0}: ".format(attr_name))
self._internal_print(attr_val, indent, file_handle)
else:
file_handle.write("- {0}: ".format(attr_name))
try:
str(attr_val).decode("ascii")
file_handle.write("{0}".format(attr_val))
except UnicodeDecodeError:
# python2 - string contains binary data
file_handle.write("<binary data>")
except AttributeError:
# python3 - does not require decoding strings and thus falls into this code
if isinstance(attr_val, bytes):
file_handle.write("<binary data>")
else:
file_handle.write("{0}".format(attr_val))
file_handle.write("\n")