graph-tool 20.8 KB
Newer Older
Tiago Peixoto's avatar
Tiago Peixoto committed
1
#! /usr/bin/env python
2
# graph-tool -- a general graph manipulation stand-alone tool
3
#
Tiago Peixoto's avatar
Tiago Peixoto committed
4
# Copyright (C) 2007 Tiago de Paula Peixoto <tiago@forked.de>
Tiago Peixoto's avatar
Tiago Peixoto committed
5
#
6
# This program is free software: you can redistribute it and/or modify
Tiago Peixoto's avatar
Tiago Peixoto committed
7
# it under the terms of the GNU General Public License as published by
8
# the Free Software Foundation, either version 3 of the License, or
Tiago Peixoto's avatar
Tiago Peixoto committed
9
# (at your option) any later version.
10
#
Tiago Peixoto's avatar
Tiago Peixoto committed
11
12
13
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
14
15
# GNU General Public License for more details.
#
Tiago Peixoto's avatar
Tiago Peixoto committed
16
# You should have received a copy of the GNU General Public License
17
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
Tiago Peixoto's avatar
Tiago Peixoto committed
18

19
from graph_tool import *
Tiago Peixoto's avatar
Tiago Peixoto committed
20
from graph_tool import libgraph_tool, _eval_expr, _open_file
21
22

import sys, os, os.path, re, struct, fcntl, termios, gzip, bz2, string,\
23
24
       textwrap, time, signal, traceback, shutil, time, math, inspect, \
       functools, types, optparse
Tiago Peixoto's avatar
Tiago Peixoto committed
25

26
27
def _main():
    """Function to be run only when module is executed"""
28

29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
    # General utilities
    def getheightwidth():
        """This will get the terminal width and height"""
        height, width = 20, 80
        try:
            height, width = struct.unpack("hhhh",
                                          fcntl.ioctl(0, termios.TIOCGWINSZ,
                                                      "\000"*8))[0:2]
        except IOError:
            pass
        return height, width

    ############################################################################
    # Command line parsing
    ############################################################################

    class OptionError(Exception):
        """Parse error in a command line option"""
        def __init__(self, option, message):
            self.option = option
            self.message = message
50
        def __str__(self):
51
52
53
54
55
56
57
58
59
            if self.option != None:
                msg = "error parsing option --%s: %s\n\n" % (self.option.name,
                                                             self.message)
                help =  "--%s=%s\n" % (self.option.name,
                                       self.option.get_metavars())
                wrapper = textwrap.TextWrapper(initial_indent = " "*3,
                                               subsequent_indent =  " "*3,
                                               width = getheightwidth()[1])
                msg += help + wrapper.fill(self.option.get_help())
60
            else:
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
                msg = self.message
            return msg

    class OptionMask:
        """This class masks a function as command line option. The idea is that
        something like:
            --foo="baz|gnat|7.0||bar=42, opt='abc'"
        should be translated to:
            foo("baz", "gnat", 7.0, None, bar=42, opt='abc')

        The argument types should be deduced from the default arguments values.
        In general, there are two types of option values: positional (separated
        by '|') and sub-options (separated by ', '). Sub-options always come
        after all other positional arguments, and are themselves, as a whole, a
        positional argument. A positional argument is considered optional if it
        has a default value in the underlying function. Suboptions are
        identified by the 'first_subopt' attribute, which, if existent, tells
        which of the function parameters correspond to the start of the
        sub-options. If the function returns a value (has a 'has_output'
        attribute), an optional 'file' option is appended to its option list, so
        it can be redirected to a file.
        """

        def __init__(self, func):
            self.func = func
            self.name = func.__name__.replace("_", "-")

            arg_spec = inspect.getargspec(func)

            if func.__class__ == types.MethodType:
                args = arg_spec[0][1:] # bound methods always have a spurious
                                       # first argument called 'self'
93
            else:
94
95
96
                args = arg_spec[0]
            defaults = arg_spec[3]
            self.args = args
97

98
99
100
101
102
            if defaults != None:
                self.defaults = dict([ (args[-i], defaults[-i]) \
                                       for i in xrange(1, len(defaults)+1) ])
            else:
                self.defaults = {}
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
            if func.func_dict.has_key("first_subopt"):
                first_subopt = func.func_dict["first_subopt"]
                self.positional = args[0:args.index(first_subopt)]
                self.subopts = args[args.index(first_subopt):]
            else:
                self.positional = args
                self.subopts = None

            self.has_output = func.func_dict.get("has_output", False)
            if self.has_output:
                self.output_file = None

        def get_metavars(self):
            """This returns the meta-variables string, ex: [foo]|bar|[baz] """
            metavars = []
            for arg in self.positional:
                if self.defaults.has_key(arg):
                    metavars.append("[%s]" % arg)
                else:
                    metavars.append(arg)
            if self.subopts != None:
                metavars.append("[options]")
            if self.has_output:
                metavars.append("[output_file]")
            return "|".join(metavars) if len(metavars) > 0 else None

        def get_help(self):
            """This returns the option's help string"""
            if self.subopts != None:
                opts = []
                for o in self.subopts:
                    if self.defaults.has_key(o):
                        if self.defaults[o].__class__ == str:
                            opts.append(o + " (default: '%s')" % \
                                        self.defaults[o])
                        else:
                            opts.append(o + " (default: %s)" % \
                                        str(self.defaults[o]))
                    else:
                        opts.append(o)
                opts = " Options are: " + ", ".join(opts)
            else:
                opts = ""
            return inspect.getdoc(self.func).replace("\n", " ")  + opts
148

149
150
151
        def parse(self, arg, dry_run=False):
            """this parses the command line argument in 'arg' and calls the
            underlying function"""
152
153
            # Parameter regexp
            p = re.compile(r"((('[^']*')|(\"[^\"']*\"))|([^\|]+)|(^(?=[^$]))(?=\||$))")
154
155
            if arg == None:
                arg = ""
156
157
158
159
160
161
162
163
164
165
166
167
168
            values = [x[0].strip() for x in p.findall(arg)]

            n_min = len([x for x in self.positional \
                         if not self.defaults.has_key(x)])
            n_max = len(self.positional)
            if self.subopts != None:
                n_max += 1
            if self.has_output:
                n_max += 1

            if len(values) > n_max or len(values) < n_min:
                if n_min != n_max:
                    expected = "between %d and %d" % (n_min, n_max)
169
                else:
170
171
172
173
174
175
176
177
178
179
180
181
182
183
                    expected = str(n_min)
                raise OptionError(self, "invalid number of arguments: %d" % \
                                  len(values) +" (expected %s)" % expected)

            # parse suboptions
            subopt_list = {}
            if self.subopts != None and len(values) > len(self.positional):
                p = re.compile(r"([^, '\"]+|[^, ]+'[^']*'|[^, ]+\"[^\"]*\")(?:, |$)")
                subopt_list = dict([ (x.split("=", 1)[0].strip(),
                                      x.split("=", 1)[1].strip()) \
                                     for x in p.findall(values[-1]) \
                                     if "=" in x])
                for k, v in subopt_list.iteritems():
                    if self.defaults.has_key(k):
184
                        try:
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
                            vtype = self.defaults[k].__class__
                            subopt_list[k] = vtype(v)
                        except ValueError, e:
                            raise OptionError(self,
                                              "invalid value for option" + \
                                              " '%s': %s (%s)" % (k, v, str(e)))
                    else:
                        raise OptionError(self, "invalid option: %s" % k)
                bool_list = dict([ (x.strip(), True) \
                                   for x in p.findall(values[-1]) \
                                   if "=" not in x])
                for k in bool_list.keys():
                    if self.defaults.has_key(k):
                        if self.defaults[k].__class__ != bool:
                            raise OptionError(self, "option '%s' " % k +
                                              "requires an argument")
                    else:
                        raise OptionError(self, "invalid option: %s" % k)
                subopt_list.update(bool_list)

            # setup args and call function
            args = values
            if len(args) - len(self.positional) == 2:
                self.output_file = args.pop()
                args.pop() # suboptions
            elif len(args) - len(self.positional) == 1:
                if self.subopts != None:
                    args.pop()
213
                else:
214
                    self.output_file = args.pop()
215

216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
            args = tuple(args)
            if not dry_run:
                if len(subopt_list) > 0:
                    return self.func(*args, **subopt_list)
                else:
                    return self.func(*args)

    def populate_options(graph, parsed_options):
        """This function will take an instance of a Graph class and populate the
        command line option parser with its methods, and will return a tuple
        with the parser, and the list of options. Later, during the actual
        parsing, the parsed_options list will keep each option which was found
        in the command line, together with its value, after being parsed"""

        # We will take 'graph', which is an instance of a Graph class, and
        # collect all its appropriate methods as command line options. They will
        # be put in a list of groups, where each group itself is a list where
        # the first element is the name and the others are the options, which
        # point to the corresponding method in the Graph class, through an
        # instance of an OptionMask class.

237
        option_groups = [ [x] for x in graph._Graph__groups ]
238
239
        functions = [ getattr(graph, f) for f in dir(graph) ]
        functions = [ f for f in functions if callable(f) and \
240
241
                      "opt_group" in dir(f) ]

242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
        for g in option_groups:
            g.extend([OptionMask(f) for f in functions \
                      if f.func_dict["opt_group"] == g[0]])

        # Setup 'optparse' and load it with all the options
        prog_info = libgraph_tool.mod_info()
        version_string = "%s %s\nWritten by %s\n\n%s" % (prog_info.name,
                                                         prog_info.version,
                                                         prog_info.author,
                                                         prog_info.copyright)

        # The option parser
        class MyOptionParser(optparse.OptionParser):
            def error(self, error):
                raise OptionError(None, error)
        (height, width) = getheightwidth()
        fmtr = optparse.IndentedHelpFormatter(max_help_position=width/2,
                                              width=width)
        parser = MyOptionParser(usage="%s [options]" % sys.argv[0],
                                version=version_string,
                                formatter=fmtr)

        # the parsed_options list will keep each option which was found in the
        # command line, together with its value, after being parsed
        def push_option(option, opt_str, value, parser, opt=None,
                        option_list=None):
            "pushes the option into the option list"
            try:
                opt.parse(value, dry_run=True)
            except OptionError, e:
                parser.error(str(e))
            option_list.append((opt, value))

        # Populate the option parser
        options = {} # maps each option by name
        for group in option_groups:
            g = parser.add_option_group(group[0])
            for option in group[1:]:
                if option.get_metavars() != None:
                    g.add_option("--%s" % option.name, action="callback",
                                 callback=push_option,
                                 callback_kwargs={"opt":option,
                                                  "option_list":parsed_options},
                                 metavar=option.get_metavars(), type="string",
                                 help=option.get_help())
                else:
                    g.add_option("--%s" % option.name, action="callback",
                                 callback=push_option,
                                 callback_kwargs={"opt":option,
                                                  "option_list":parsed_options},
                                 help=option.get_help())
                options[option.name] = option

        # add --for option
        g = parser.add_option_group("Scripting")
        g.add_option("--for", dest="for_arg",
                     metavar="init_expr|cond_expr|step_expr|[prefix_expr]",
                     type="string", help="simplified scripting")

        # build option list
        options = []
        for g in option_groups:
            for o in g[1:]:
                options.append(o)

        return (parser, options)

    def print_result(retval, file=sys.stdout, prefix=""):
        """This will print a given result, according to its type. It will also
        prefix each line with 'prefix'"""
        if retval.__class__ == dict:
            keys = retval.keys()
            keys.sort()
            for k in keys:
                if k.__class__ == tuple or k.__class__ == list:
                    print >> file, "%s%s" % (prefix, " \t".join(["%.20g" % x \
                                                                 for x in k])),
                else:
                    print >> file, "%s%.20g" % (prefix, k),
                if retval[k].__class__ == tuple or retval[k].__class__ == list:
                    print >> file, " \t", " \t".join(["%.20g" % x for x\
                                                      in retval[k]])
                else:
                    print >> file, " \t%.20g" % retval[k]
326
        else:
327
328
329
            if retval.__class__ == list or retval.__class__ == tuple:
                print >> file, "%s%s" % (prefix, " \t".join(["%.20g" % x \
                                                             for x in retval]))
Tiago Peixoto's avatar
Tiago Peixoto committed
330
            else:
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
                print >> file, "%s%.20g" % (prefix, retval)
        if hasattr(file, 'flush'):
            file.flush()

    class RunError(Exception):
        """Error raised during option run"""
        def __init__(self, option, value, message):
            self.option = option
            self.value = value
            self.message = message
        def __str__(self):
            return "error during evaluation of command --%s=%s: %s" % \
                   (self.option.name, self.value, self.message)

    def run(option_list, output_prefix=""):
        """This will parse the options and run the specified actions."""
        for opt in option_list:
            # Run the option itself, and report on any error
            try:
                retval = opt[0].parse(opt[1])
            except (GraphError, OptionError, IOError, TypeError, ValueError), e:
                raise RunError(opt[0], opt[1], str(e))

            # if there is a return value, store it or display it properly
            if retval != None:
                output_file = opt[0].output_file
                if output_file == "-" or output_file == None:
                    if opt[1] != "":
                        print "# %s (%s)" % (opt[0].name.replace("-", " "),
                                             opt[1])
361
                    else:
362
                        print "# %s" % opt[0].name.replace("-", " ")
363
                    print_result(retval, prefix=output_prefix)
364
                    print
Tiago Peixoto's avatar
Tiago Peixoto committed
365
                else:
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
                    print_result(retval,
                                 _open_file(output_file,
                                            mode = "w" if output_prefix == "" \
                                                    else "a"),
                                 prefix=output_prefix)

    def run_loop(expr, arguments, parsed_options):
        """This will loop through the loop specification given by the arguments,
        parse the options and run the specified actions."""

        class LoopTemplate(string.Template):
            delimiter = "%"
        try:
            p = re.compile(r"((('[^']*')|(\"[^\"']*\"))|([^\|]+)|(^(?=[^$]))(?=\||$))")
            loop_args = [x[0].strip() for x in p.findall(expr)]
            if len(loop_args) not in range(3,5):
                raise RunError("for", expr, "invalid loop specification")

            loop_vars = dict()
            _eval_expr(loop_args[0] + "; None", loop_vars)
            cond_expr = loop_args[1]
            step_expr = loop_args[2]
            if len(loop_args) > 3:
                prefix_expr = loop_args[3]
            else:
                prefix_expr = "''"
            while eval(cond_expr, loop_vars):
                args = []
                for arg in arguments:
                    try:
                        args.append(LoopTemplate(arg).substitute(loop_vars))
                    except KeyError, e:
                        raise RunError("for", expr,
                                       "key %s not found at '%s'" % (e, arg))
                prefix = eval(prefix_expr, loop_vars)
                if prefix.__class__ != str:
                    prefix = "%.20g\t " % prefix

                # parse the command line
                while len(parsed_options) > 0:
                    parsed_options.pop()
                (opt_val, pargs) =  parser.parse_args(args)
                if len(pargs) > 0:
                    parser.error("invalid arguments: " + " ".join(pargs))

                run(parsed_options, prefix)

                exec step_expr in loop_vars
        except SyntaxError, e:
            raise RunError("for", expr, "error evaluating loop: " % e)

    # Do the actual parsing.
    try:
        g = Graph() # the one and only graph

        # initialize internal signal handling
        g._Graph__graph.InitSignalHandling()

        parsed_options = []
        (parser, option_list) = populate_options(g, parsed_options)

        # A hack to allow optional option arguments. We'll just append a '=' to
        # options with arguments and have been supplied with zero arguments.
        arguments = sys.argv[1:]
        args = []
        opt_names = [o.name for o in option_list]
        for i in xrange(0, len(arguments)):
            a = arguments[i]
            if i < len(arguments) - 1:
                n = arguments[i+1] # next parameter
            else:
                n = a
            if a[2:] in opt_names:
                option = option_list[opt_names.index(a[2:])]
                if (option.get_metavars() != None or option.has_output) and \
                       n.startswith("--"):
                    args.append(a + "=")
                else:
                    args.append(a)
445
            else:
446
                args.append(a)
447

448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
        # Support bashcompletion, with the optcomplete module. See
        # http://furius.ca/optcomplete for details.
        try:
            import optcomplete
            optcomplete.autocomplete(parser)
        except ImportError:
            pass

        # separate "normal" args from args after the "--for" option
        normal_args, for_args = [], []
        arg_list = normal_args
        for arg in args:
            arg_list.append(arg)
            if arg.startswith("--for"):
                arg_list = for_args

        # parse the command line
        (opt_val, pargs) = parser.parse_args(normal_args)
        if len(pargs) > 0:
            parser.error("invalid arguments: " + " ".join(pargs))

        # execute the options
        run(parsed_options)

        if opt_val.for_arg != None:
            run_loop(opt_val.for_arg, for_args, parsed_options)

    except (OptionError, RunError), e:
476
477
478
479
480
481
482
483
484
485
        print >> sys.stderr, "graph-tool:", e
        sys.exit(1)
    except SystemExit, e:
        sys.exit(e)
    except KeyboardInterrupt:
        print >> sys.stderr, "graph-tool error: Keyboard interrupt."
        sys.exit(3)
    except MemoryError, e:
        error = str(e)
        if error != "":
486
487
            print >> sys.stderr, "graph-tool error: out of memory. (text:",
            print >> sys.stderr, error, ")"
488
489
490
491
492
493
        else:
            print >> sys.stderr, "graph-tool error: out of memory."
        sys.exit(4)

# when executed, run main():
if __name__ == '__main__':
494
    _main()