OLD | NEW |
| (Empty) |
1 #!/usr/bin/env python | |
2 # Copyright (c) 2011 The Chromium Authors. All rights reserved. | |
3 # Use of this source code is governed by a BSD-style license that can be | |
4 # found in the LICENSE file. | |
5 | |
6 """Generates profile dictionaries for Autofill. | |
7 | |
8 Used to test autofill.AutofillTest.FormFillLatencyAfterSubmit. | |
9 Can be used as a stand alone script with -h to print out help text by running: | |
10 python autofill_dataset_generator.py -h | |
11 """ | |
12 | |
13 import codecs | |
14 import logging | |
15 from optparse import OptionParser | |
16 import os | |
17 import random | |
18 import re | |
19 import sys | |
20 | |
21 | |
22 class NullHandler(logging.Handler): | |
23 def emit(self, record): | |
24 pass | |
25 | |
26 | |
27 class DatasetGenerator(object): | |
28 """Generates a dataset of dictionaries. | |
29 | |
30 The lists (such as address_construct, city_construct) define the way the | |
31 corresponding field is generated. They accomplish this by specifying a | |
32 list of function-args lists. | |
33 """ | |
34 address_construct = [ | |
35 [ random.randint, 1, 10000], | |
36 [ None, u'foobar'], | |
37 [ random.choice, [ u'St', u'Ave', u'Ln', u'Ct', ]], | |
38 [ random.choice, [ u'#1', u'#2', u'#3', ]], | |
39 ] | |
40 | |
41 city_construct = [ | |
42 [ random.choice, [ u'San Jose', u'San Francisco', u'Sacramento', | |
43 u'Los Angeles', ]], | |
44 ] | |
45 | |
46 state_construct = [ | |
47 [ None, u'CA'] | |
48 ] | |
49 | |
50 # These zip codes are now matched to the corresponding cities in | |
51 # city_construct. | |
52 zip_construct = [ u'95110', u'94109', u'94203', u'90120'] | |
53 | |
54 logger = logging.getLogger(__name__) | |
55 logger.addHandler(NullHandler()) | |
56 log_handlers = {'StreamHandler': None} | |
57 | |
58 def __init__(self, output_filename=None, logging_level=None): | |
59 """Constructs dataset generator object. | |
60 | |
61 Creates 'fields' data member which is a list of pair (two values) lists. | |
62 These pairs are comprised of a field key e.g. u'NAME_FIRST' and a | |
63 generator method e.g. self.GenerateNameFirst which will generate the value. | |
64 If we want the value to always be the same e.g. u'John' we can use this | |
65 instead of a method. We can even use None keyword which will give | |
66 a value of u''. | |
67 | |
68 'output_pattern' for one field would have been: "{u'NAME_FIRST': u'%s',}" | |
69 which is ready to accept a value for the 'NAME_FIRST' field key once | |
70 this value is generated. | |
71 'output_pattern' is used in 'GenerateNextDict()' to generate the next | |
72 dict line. | |
73 | |
74 Args: | |
75 output_filename: specified filename of generated dataset to be saved. | |
76 Default value is None and no saving takes place. | |
77 logging_level: set verbosity levels, default is None. | |
78 """ | |
79 if logging_level: | |
80 if not self.log_handlers['StreamHandler']: | |
81 console = logging.StreamHandler() | |
82 console.setLevel(logging.INFO) | |
83 self.log_handlers['StreamHandler'] = console | |
84 self.logger.addHandler(console) | |
85 self.logger.setLevel(logging_level) | |
86 else: | |
87 if self.log_handlers['StreamHandler']: | |
88 self.logger.removeHandler(self.log_handlers['StreamHandler']) | |
89 self.log_handlers['StreamHandler'] = None | |
90 | |
91 self.output_filename = output_filename | |
92 | |
93 self.dict_no = 0 | |
94 self.fields = [ | |
95 [u'NAME_FIRST', self.GenerateNameFirst], | |
96 [u'NAME_MIDDLE', None], | |
97 [u'NAME_LAST', None], | |
98 [u'EMAIL_ADDRESS', self.GenerateEmail], | |
99 [u'COMPANY_NAME', None], | |
100 [u'ADDRESS_HOME_LINE1', self.GenerateAddress], | |
101 [u'ADDRESS_HOME_LINE2', None], | |
102 [u'ADDRESS_HOME_CITY', self.GenerateCity], | |
103 [u'ADDRESS_HOME_STATE', self.GenerateState], | |
104 [u'ADDRESS_HOME_ZIP', self.GenerateZip], | |
105 [u'ADDRESS_HOME_COUNTRY', u'United States'], | |
106 [u'PHONE_HOME_WHOLE_NUMBER', None], | |
107 ] | |
108 | |
109 self.next_dict = {} | |
110 # Using implicit line joining does not work well in this case as each line | |
111 # has to be strings and not function calls that may return strings. | |
112 self.output_pattern = u'{\'' + \ | |
113 u', '.join([u'u"%s" : u"%%s"' % key for key, method in self.fields]) + \ | |
114 u',}' | |
115 | |
116 def _GenerateField(self, field_construct): | |
117 """Generates each field in each dictionary. | |
118 | |
119 Args: | |
120 field_construct: it is a list of lists. | |
121 The first value (index 0) of each containing list is a function or None. | |
122 The remaining values are the args. If function is None then arg is just | |
123 returned. | |
124 | |
125 Example 1: zip_construct = [[ None, u'95110']]. There is one | |
126 containing list only and function here is None and arg is u'95110'. | |
127 This just returns u'95110'. | |
128 | |
129 Example 2: address_construct = [ [ random.randint, 1, 10000], | |
130 [ None, u'foobar'] ] This has two containing lists and it will return | |
131 the result of: | |
132 random.randint(1, 10000) + ' ' + u'foobar' | |
133 which could be u'7832 foobar' | |
134 """ | |
135 parts = [] | |
136 for function_and_args in field_construct: | |
137 function = function_and_args[0] | |
138 args = function_and_args[1:] | |
139 if not function: | |
140 function = lambda x: x | |
141 parts.append(str(function(*args))) | |
142 return (' ').join(parts) | |
143 | |
144 def GenerateAddress(self): | |
145 """Uses _GenerateField() and address_construct to gen a random address. | |
146 | |
147 Returns: | |
148 A random address. | |
149 """ | |
150 return self._GenerateField(self.address_construct) | |
151 | |
152 def GenerateCity(self): | |
153 """Uses _GenerateField() and city_construct to gen a random city. | |
154 | |
155 Returns: | |
156 A random city. | |
157 """ | |
158 return self._GenerateField(self.city_construct) | |
159 | |
160 def GenerateState(self): | |
161 """Uses _GenerateField() and state_construct to generate a state. | |
162 | |
163 Returns: | |
164 A state. | |
165 """ | |
166 return self._GenerateField(self.state_construct) | |
167 | |
168 def GenerateZip(self): | |
169 """Uses zip_construct and generated cities to return a matched zip code. | |
170 | |
171 Returns: | |
172 A zip code matched to the corresponding city. | |
173 """ | |
174 city_selected = self.next_dict['ADDRESS_HOME_CITY'][0] | |
175 index = self.city_construct[0][1].index(city_selected) | |
176 return self.zip_construct[index] | |
177 | |
178 def GenerateCountry(self): | |
179 """Uses _GenerateField() and country_construct to generate a country. | |
180 | |
181 Returns: | |
182 A country. | |
183 """ | |
184 return self._GenerateField(self.country_construct) | |
185 | |
186 def GenerateNameFirst(self): | |
187 """Generates a numerical first name. | |
188 | |
189 The name is the number of the current dict. | |
190 i.e. u'1', u'2', u'3' | |
191 | |
192 Returns: | |
193 A numerical first name. | |
194 """ | |
195 return u'%s' % self.dict_no | |
196 | |
197 def GenerateEmail(self): | |
198 """Generates an email that corresponds to the first name. | |
199 | |
200 i.e. u'1@example.com', u'2@example.com', u'3@example.com' | |
201 | |
202 Returns: | |
203 An email address that corresponds to the first name. | |
204 """ | |
205 return u'%s@example.com' % self.dict_no | |
206 | |
207 | |
208 def GenerateNextDict(self): | |
209 """Generates next dictionary of the dataset. | |
210 | |
211 Returns: | |
212 The output dictionary. | |
213 """ | |
214 self.dict_no += 1 | |
215 self.next_dict = {} | |
216 for key, method_or_value in self.fields: | |
217 if not method_or_value: | |
218 self.next_dict[key] = [''] | |
219 elif type(method_or_value) in [str, unicode]: | |
220 self.next_dict[key] = ['%s' % method_or_value] | |
221 else: | |
222 self.next_dict[key] = [method_or_value()] | |
223 return self.next_dict | |
224 | |
225 def GenerateDataset(self, num_of_dict_to_generate=10): | |
226 """Generates a list of dictionaries. | |
227 | |
228 Args: | |
229 num_of_dict_to_generate: The number of dictionaries to be generated. | |
230 Default value is 10. | |
231 | |
232 Returns: | |
233 The dictionary list. | |
234 """ | |
235 random.seed(0) # All randomly generated values are reproducible. | |
236 if self.output_filename: | |
237 output_file = codecs.open( | |
238 self.output_filename, mode='wb', encoding='utf-8-sig') | |
239 else: | |
240 output_file = None | |
241 try: | |
242 list_of_dict = [] | |
243 if output_file: | |
244 output_file.write('[') | |
245 output_file.write(os.linesep) | |
246 | |
247 while self.dict_no < num_of_dict_to_generate: | |
248 output_dict = self.GenerateNextDict() | |
249 list_of_dict.append(output_dict) | |
250 output_line = self.output_pattern % tuple( | |
251 [output_dict[key] for key, method in self.fields]) | |
252 if output_file: | |
253 output_file.write(output_line) | |
254 output_file.write(os.linesep) | |
255 self.logger.info( | |
256 '%d: [%s]' % (self.dict_no, output_line.encode(sys.stdout.encoding, | |
257 'ignore'))) | |
258 | |
259 if output_file: | |
260 output_file.write(']') | |
261 output_file.write(os.linesep) | |
262 self.logger.info('%d dictionaries generated SUCCESSFULLY!', self.dict_no) | |
263 self.logger.info('--- FINISHED ---') | |
264 return list_of_dict | |
265 finally: | |
266 if output_file: | |
267 output_file.close() | |
268 | |
269 | |
270 def main(): | |
271 parser = OptionParser() | |
272 parser.add_option( | |
273 '-o', '--output', dest='output_filename', default='', | |
274 help='write output to FILE [optional]', metavar='FILE') | |
275 parser.add_option( | |
276 '-d', '--dict', type='int', dest='dict_no', metavar='DICT_NO', default=10, | |
277 help='DICT_NO: number of dictionaries to be generated [default: %default]') | |
278 parser.add_option( | |
279 '-l', '--log_level', dest='log_level', default='debug', | |
280 metavar='LOG_LEVEL', | |
281 help='LOG_LEVEL: "debug", "info", "warning" or "error" [default: %default]') | |
282 | |
283 (options, args) = parser.parse_args() | |
284 if args: | |
285 parser.print_help() | |
286 return 1 | |
287 options.log_level = options.log_level.lower() | |
288 if options.log_level not in ['debug', 'info', 'warning', 'error']: | |
289 parser.error('Wrong log_level argument.') | |
290 parser.print_help() | |
291 else: | |
292 if options.log_level == 'debug': | |
293 options.log_level = logging.DEBUG | |
294 elif options.log_level == 'info': | |
295 options.log_level = logging.INFO | |
296 elif options.log_level == 'warning': | |
297 options.log_level = logging.WARNING | |
298 elif options.log_level == 'error': | |
299 options.log_level = logging.ERROR | |
300 | |
301 gen = DatasetGenerator(options.output_filename, options.log_level) | |
302 gen.GenerateDataset(options.dict_no) | |
303 return 0 | |
304 | |
305 | |
306 if __name__ == '__main__': | |
307 sys.exit(main()) | |
OLD | NEW |