core.py 33.4 KB
Newer Older
1
#! /usr/bin/env python
2
# graph_tool -- a general graph manipulation python module
3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18
#
# Copyright (C) 2007 Tiago de Paula Peixoto <tiago@forked.de>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# 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
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

Tiago Peixoto's avatar
Tiago Peixoto committed
19 20
from dl_import import *
dl_import("import libgraph_tool_core as libcore")
21
__version__ = libcore.mod_info().version
22

23
import io # sets up libcore io routines
24

Tiago Peixoto's avatar
Tiago Peixoto committed
25
import sys, os, os.path, re, struct, fcntl, termios, gzip, bz2, string,\
26
       textwrap, time, signal, traceback, shutil, time, math, inspect, \
27
       functools, types, weakref, copy
28
from StringIO import StringIO
29
from decorators import _wraps, _require, _attrs, _handle_exceptions, _limit_args
30 31 32 33 34

################################################################################
# Utility functions
################################################################################

35 36 37 38 39 40 41
def _prop(t, g, prop):
    """Returns either a property map, or an internal vertex property map with a
    given name"""
    if type(prop) == str:
        try:
            pmap = g.properties[(t,prop)]
        except KeyError:
42
            raise GraphError("no internal %s property named: %s" %\
43 44 45 46
                             ("vertex" if t == "v" else \
                              ("edge" if t == "e" else "graph"),prop))
    else:
        pmap = prop
47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64
    if pmap == None:
        return libcore.any()
    else:
        return pmap._PropertyMap__map.get_map()

def _degree(g, name):
    """Retrieve the degree type from string, or returns the corresponding
    property map"""
    deg = name
    if name == "in-degree" or name == "in":
        deg = libcore.Degree.In
    elif name == "out-degree" or name == "out":
        deg = libcore.Degree.Out
    elif name == "total-degree" or name == "total":
        deg = libcore.Degree.Total
    else:
        deg = _prop("v", g, deg)
    return deg
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
def _parse_range(range):
    """Parse a range in the form 'lower[*] upper[*]' (where an '*' indicates
    a closed boundary) or 'comp val', where 'comp' is [==|=|!=|>=|<=|>|<].
    """
    try:
        rcomp = re.compile(r"(>|<|>=|<=|==|=|!=)\s*([^=<>!]+)")
        rrange = re.compile(r"([^\*]+)(\*?)\s+([^\*]+)(\*?)")
        inverted = False
        if rcomp.match(range):
            comp = rcomp.match(range).group(1)
            thres = rcomp.match(range).group(2)
            thres = float(thres)
            if "<" in comp:
                ran = (float('-Inf'), thres)
                if "=" in comp:
                    inc = (False, True)
                else:
                    inc = (False, False)
            elif ">" in comp:
                ran = (thres, float('Inf'))
                if "=" in comp:
                    inc = (True, False)
                else:
                    inc = (False, False)
            elif comp == "==" or comp == "=" or comp == "!=":
                ran = (thres, thres)
                inc = (True, True)
                if comp == "!=":
                    inverted = True
        elif rrange.match(range):
            m = rrange.match(range)
            r1 = float(m.group(1))
            i1 = "*" in m.group(2)
            r2 = float(m.group(3))
            i2 = "*" in m.group(4)
            ran = (r1, r2)
            inc = (i1, i2)
        else:
            raise ValueError("invalid value for range: " + range)
    except (ValueError, TypeError), e:
        raise ValueError("invalid value for range: %s: %s " % (range, str(e)))
    return (ran, inc, inverted)

109 110 111 112 113 114 115 116 117 118 119
def _type_alias(type_name):
    alias = {"int8_t":"bool",
             "boolean":"bool",
             "int":"int32_t",
             "long":"int32_t",
             "long long":"int64_t",
             "object":"python::object"}
    if type_name in value_types():
        return type_name
    if alias.has_key(type_name):
        return alias[type_name]
120
    ma = re.compile(r"vector<(.*)>").match(type_name)
121 122 123 124 125 126
    if ma:
        t = ma.group(1)
        if alias.has_key(t):
            return "vector<%s>" % alias[t]
    raise ValueError("invalid property value type: " + type_name)

Tiago Peixoto's avatar
Tiago Peixoto committed
127 128 129 130 131 132 133 134 135 136
def show_config():
    info = libcore.mod_info()
    print "version:", info.version
    print "gcc version:", info.gcc_version
    print "compilation flags:", info.cxxflags
    print "install prefix:", info.install_prefix
    print "python dir:", info.python_dir
    print "graph filtering:", libcore.graph_filtering_enabled()
    print "openmp:", libcore.openmp_enabled()
    print "uname:", " ".join(os.uname())
137

138
################################################################################
139
# Property Maps
140 141
################################################################################

142
class PropertyMap(object):
143
    """Property map class"""
144 145
    def __init__(self, pmap, g, key_type, key_trans = None):
        self.__map = pmap
146
        self.__g = weakref.ref(g)
147 148 149 150 151
        self.__key_type = key_type
        self.__key_trans = key_trans if key_trans != None else lambda k: k
        self.__register_map()

    def __register_map(self):
152 153 154 155
        if self.__g() == None:
            return
        self.__g()._Graph__known_properties.append((self.key_type(),
                                                    weakref.ref(self.__map)))
156
    def __unregister_map(self):
157 158 159 160 161
        if self.__g() == None:
            return
        i = self.__g()._Graph__known_properties.index((self.key_type(),
                                                       weakref.ref(self.__map)))
        del self.__g()._Graph__known_properties[i]
162 163 164 165 166 167 168 169 170 171 172

    def __del__(self):
        self.__unregister_map()

    def __getitem__(self, k):
        return self.__map[self.__key_trans(k)]

    def __setitem__(self, k, v):
        self.__map[self.__key_trans(k)] = v

    def get_graph(self):
173
        """Get the graph class to which the map refers"""
174
        return self.__g()
175 176 177 178 179 180 181 182 183

    def key_type(self):
        """The key type of the map. Either 'g', 'v' or 'e'"""
        return self.__key_type

    def value_type(self):
        """The value type of the map"""
        return self.__map.value_type()

184 185 186 187 188 189 190
    def get_array(self):
        """Get an array with property values

        .. WARNING::

           The returned array does not own the data, which belongs to the
           property map. Therefore, the returned array cannot have a longer
Tiago Peixoto's avatar
Tiago Peixoto committed
191
           lifespan than the property map itself! Furthermore, if the graph
192 193 194 195 196 197 198 199 200 201 202 203
           changes, it may leave the pointer to the data in the array dangling!
           Do *not* store the array if the graph is to be modified, or the
           original property map deleted; *store a copy instead*!
        """
        if self.__key_type == 'v':
            n = self.__g().num_vertices()
        elif self.__key_type == 'e':
            n = self.__g().num_edges()
        else:
            n = 1
        return self.__map.get_array(n)

204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224
class PropertyDict(dict):
    """Wrapper for the dict of vertex, graph or edge properties, which sets the
    value on the property map when changed in the dict."""
    def __init__(self, g, old, get_func, set_func, del_func):
        dict.__init__(self)
        dict.update(self, old)
        self.g = g
        self.get_func = get_func
        self.set_func = set_func
        self.del_func = del_func

    def __setitem__(self, key, val):
        if self.set_func != None:
            self.set_func(self.g, key, val)
        else:
            raise KeyError("Property dict cannot be set")
        dict.__setitem__(self, key, val)

    def __delitem__(self, key):
        self.del_func(self.g, key)
        dict.__delitem__(self, key)
225 226 227 228 229 230

################################################################################
# Graph class
# The main graph interface
################################################################################

231 232 233 234
from libgraph_tool_core import Vertex, Edge, GraphError,\
     Vector_bool, Vector_int32_t,  Vector_int64_t, Vector_double,\
     Vector_long_double, Vector_string, new_vertex_property, new_edge_property,\
     new_graph_property
235 236

class Graph(object):
237 238 239
    """This class encapsulates either a directed multigraph (default or if
    ``directed=True``) or an undirected multigraph (if ``directed=False``), with
    optional internal edge, vertex or graph properties.
240

241
    It is implemented as an adjacency list, where both vertex and edge lists are
Tiago Peixoto's avatar
Tiago Peixoto committed
242
    C++ STL vectors.
243
    """
244

245 246 247 248
    def __init__(self, g = None, directed=True):
        """Construct a graph. If 'g' is specified, the graph (and its internal
        properties) will be copied. The 'directed' parameter specified whether
        the graph should be directed or undirected."""
249 250 251 252 253 254 255 256 257
        self.__properties = {}
        self.__known_properties = []
        self.__filter_state = {"reversed": False,
                               "edge_filter": (None,False),
                               "vertex_filter": (None,False),
                               "directed": True,
                               "reversed": False}
        self.__stashed_filter_state = []

258
        if g == None:
259
            self.__graph = libcore.GraphInterface()
260
            self.set_directed(directed)
261
        else:
262
            self.__graph = libcore.GraphInterface(g.__graph)
263 264 265
            for k,v in g.__properties.iteritems():
                new_p = self.new_property(v.key_type(), v.value_type())
                self.copy_property(v, new_p, g)
266
                self.properties[k] = new_p
267
            self.__stashed_filter_state = [g.__filter_state]
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
            v_filt, v_rev = g.__filter_state["vertex_filter"]
            if v_filt != None:
                if v_filt not in g.vertex_properties.values():
                    new_filt = self.new_vertex_property("bool")
                    self.copy_property(v_filt, new_filt)

                else:
                    for k,v in g.vertex_properties.iteritems():
                        if v == v_filt:
                            new_filt = self.vertex_properties[k]
                self.__stashed_filter_state[0]["vertex_filter"] = (new_filt,
                                                                   v_rev)
            e_filt, e_rev = g.__filter_state["edge_filter"]
            if e_filt != None:
                if e_filt not in g.edge_properties.values():
                    new_filt = self.new_edge_property("bool")
                    self.copy_property(e_filt, new_filt)

                else:
                    for k,v in g.edge_properties.iteritems():
                        if v == e_filt:
                            new_filt = self.edge_properties[k]
                self.__stashed_filter_state[0]["edge_filter"] = (new_filt,
                                                                 e_rev)
            self.pop_filter()
293 294 295 296 297
        # internal index maps
        self.__vertex_index = \
                 PropertyMap(libcore.get_vertex_index(self.__graph), self, "v")
        self.__edge_index = \
                 PropertyMap(libcore.get_edge_index(self.__graph), self, "e")
298 299 300

    @_handle_exceptions
    def copy(self):
301 302 303
        """Returns a deep copy of self. All internal property maps are also
        copied."""
        return Graph(self)
304 305

    # Graph access
306
    # ============
307 308 309

    @_handle_exceptions
    def vertices(self):
310 311 312 313 314 315 316 317 318 319 320 321 322 323
        """Return an iterator over the vertices

        Examples
        --------

        >>> g = gt.Graph()
        >>> vlist = g.add_vertex(5)
        >>> vlist2 = []
        >>> for v in g.vertices():
        ...     vlist2.append(v)
        ...
        >>> assert(vlist == vlist2)

        """
324 325 326 327
        return self.__graph.Vertices()

    @_handle_exceptions
    def vertex(self, i):
328
        """Return the i-th vertex from the graph."""
329
        return self.__graph.Vertex(int(i))
330 331 332

    @_handle_exceptions
    def edges(self):
333
        """Return iterator over the edges."""
334 335 336
        return self.__graph.Edges()

    @_handle_exceptions
337 338 339 340 341 342 343
    def add_vertex(self, n=1):
        """Add a vertices to the graph, and return it. If n > 1, n vertices are
        inserted and a list is returned."""
        if n == 1:
            return self.__graph.AddVertex()
        else:
            return [self.__graph.AddVertex() for i in xrange(0,n)]
344 345

    @_handle_exceptions
346
    def remove_vertex(self, vertex):
347
        """Remove a vertex from the graph."""
348
        k = vertex.in_degree() + vertex.out_degree()
349 350
        index = self.vertex_index[vertex]
        for pmap in self.__known_properties:
351 352 353
            if pmap[0] == "v" and pmap[1]() != None and \
                   pmap[1]() != self.__vertex_index._PropertyMap__map:
                self.__graph.ShiftVertexProperty(pmap[1]().get_map(), index)
354
        self.clear_vertex(vertex)
355 356
        self.__graph.RemoveVertex(vertex)

357 358 359 360 361 362 363 364 365 366
    @_handle_exceptions
    def remove_vertex_if(self, predicate):
        """Remove all the vertices from the graph for which predicate(v)
        evaluates to True."""
        N = self.num_vertices()
        for i in xrange(0, N):
            v = self.vertex(N - i - 1)
            if predicate(v):
                self.remove_vertex(v)

367 368 369 370 371 372 373 374 375
    @_handle_exceptions
    def clear_vertex(self, vertex):
        """Removes all in and out-edges from the given vertex."""
        del_es = []
        for e in vertex.all_edges():
            del_es.append(e)
        for e in del_es:
            self.remove_edge(e)

376 377 378
    @_handle_exceptions
    def add_edge(self, source, target):
        """Add a new edge from 'source' to 'target' to the graph, and return
379
        it."""
380 381 382
        return self.__graph.AddEdge(source, target)

    @_handle_exceptions
383
    def remove_edge(self, edge):
384
        """Remove an edge from the graph."""
385 386
        self.__graph.RemoveEdge(edge)

387 388 389 390 391 392 393 394 395 396 397 398
    @_handle_exceptions
    def remove_edge_if(self, predicate):
        """Remove all the edges from the graph for which predicate(e) evaluates
        to True."""
        for v in self.vertices():
            del_es = []
            for e in v.out_edges():
                if predicate(e):
                    del_es.append(e)
            for e in del_es:
                self.remove_edge(e)

399 400
    @_handle_exceptions
    def clear(self):
401
        """Remove all vertices and edges from the graph."""
402 403
        self.__graph.Clear()

404 405
    @_handle_exceptions
    def clear_edges(self):
406
        """Remove all edges from the graph."""
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
        self.__graph.ClearEdges()

    # Internal property maps
    # ======================

    # all properties
    @_handle_exceptions
    def __get_properties(self):
        return PropertyDict(self, self.__properties,
                            lambda g,k: g.__properties[k],
                            lambda g,k,v: g.__set_property(k[0],k[1],v),
                            lambda g,k: g.__del_property(k[0],k[1]))

    @_handle_exceptions
    @_limit_args({"t":["v", "e", "g"]})
    @_require("k", str)
    @_require("v", PropertyMap)
    def __set_property(self, t, k, v):
        if t != v.key_type():
            raise ValueError("wrong key type for property map")
        self.__properties[(t,k)] = v

    @_handle_exceptions
    @_limit_args({"t":["v", "e", "g"]})
    @_require("k", str)
    def __del_property(self, t, k):
        del self.__properties[(t,k)]

    properties = property(__get_properties,
436 437 438 439 440 441 442 443 444 445 446 447
                          doc=
    """Dictionary of internal properties. Keys must always be a tuple, where the
    first element if a string from the set {'v', 'e', 'g'}, representing a
    vertex, edge or graph property, and the second element is the name of the
    property map.

    Examples
    --------
    >>> g = gt.Graph()
    >>> g.properties[("e", "foo")] = g.new_edge_property("vector<double>")
    >>> del g.properties[("e", "foo")]
    """)
448 449 450 451 452 453 454

    def __get_specific_properties(self, t):
        props = dict([(k[1],v) for k,v in self.__properties.iteritems() \
                      if k[0] == t ])
        return props

    # vertex properties
455 456
    @_handle_exceptions
    def __get_vertex_properties(self):
457 458 459 460
        return PropertyDict(self, self.__get_specific_properties("v"),
                            lambda g,k: g.__properties[("v",k)],
                            lambda g,k,v: g.__set_property("v",k,v),
                            lambda g,k: g.__del_property("v",k))
461 462
    vertex_properties = property(__get_vertex_properties,
                                 doc="Dictionary of vertex properties")
463
    # edge properties
464 465
    @_handle_exceptions
    def __get_edge_properties(self):
466 467 468 469
        return PropertyDict(self, self.__get_specific_properties("e"),
                            lambda g,k: g.__properties[("e",k)],
                            lambda g,k,v: g.__set_property("e",k,v),
                            lambda g,k: g.__del_property("e",k))
470 471 472
    edge_properties = property(__get_edge_properties,
                                 doc="Dictionary of edge properties")

473
    # graph properties
474 475
    @_handle_exceptions
    def __get_graph_properties(self):
476 477 478 479 480 481 482 483 484
        return PropertyDict(self, self.__get_specific_properties("g"),
                            lambda g,k: g.__properties[("g",k)],
                            lambda g,k,v: g.__set_property("g",k,v),
                            lambda g,k: g.__del_property("g",k))
    graph_properties = property(__get_graph_properties,
                                 doc="Dictionary of graph properties")

    @_handle_exceptions
    def list_properties(self):
485 486 487 488 489 490 491 492
        """List all internal properties for convenience.

        Examples
        --------
        >>> g = gt.Graph()
        >>> g.properties[("e", "foo")] = g.new_edge_property("vector<double>")
        >>> g.vertex_properties["foo"] = g.new_vertex_property("double")
        >>> g.vertex_properties["bar"] = g.new_vertex_property("python::object")
493 494
        >>> g.graph_properties["gnat"] = g.new_graph_property("string",\
                                                              "hi there!")
495 496 497 498 499 500 501
        >>> g.list_properties()
        gnat           (graph)   (type: string, val: hi there!)
        bar            (vertex)  (type: python::object)
        foo            (vertex)  (type: double)
        foo            (edge)    (type: vector<double>)
        """

502 503 504 505
        if len(self.__properties) == 0:
            return
        w = max([len(x[0]) for x in self.__properties.keys()]) + 4
        w = w if w > 14 else 14
506

507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522
        for k,v in self.__properties.iteritems():
            if k[0] == "g":
                print "%%-%ds (graph)   (type: %%s, val: %%s)" % w % \
                      (k[1], v.value_type(), str(v[self]))
        for k,v in self.__properties.iteritems():
            if k[0] == "v":
                print "%%-%ds (vertex)  (type: %%s)" % w % (k[1],
                                                            v.value_type())
        for k,v in self.__properties.iteritems():
            if k[0] == "e":
                print "%%-%ds (edge)    (type: %%s)" % w % (k[1],
                                                            v.value_type())

    # index properties

    def _get_vertex_index(self):
523
        return self.__vertex_index
524 525
    vertex_index = property(_get_vertex_index,
                            doc="Vertex index map. This map is immutable.")
526 527

    def _get_edge_index(self):
528
        return self.__edge_index
529
    edge_index = property(_get_edge_index, doc="Edge index map.")
530 531 532 533 534 535

    # Property map creation

    @_handle_exceptions
    def new_property(self, key_type, type):
        """Create a new (uninitialized) vertex property map of key type
536 537
        `key_type` (``v``, ``e`` or ``g``), value type ``type``, and return it.
        """
538 539 540 541 542 543 544 545 546 547
        if key_type == "v" or key_type == "vertex":
            return self.new_vertex_property(type)
        if key_type == "e" or key_type == "edge":
            return self.new_edge_property(type)
        if key_type == "g" or key_type == "graph":
            return self.new_graph_property(type)
        raise GraphError("unknown key type: " + key_type)

    @_handle_exceptions
    def new_vertex_property(self, type):
548 549
        """Create a new (uninitialized) vertex property map of type `type`, and
        return it."""
550
        return PropertyMap(new_vertex_property(_type_alias(type),
551 552
                                               self.__graph.GetVertexIndex()),
                           self, "v")
553

554 555
    @_handle_exceptions
    def new_edge_property(self, type):
556 557
        """Create a new (uninitialized) edge property map of type `type`, and
        return it."""
558 559
        return PropertyMap(new_edge_property(_type_alias(type),
                                             self.__graph.GetEdgeIndex()),
560 561 562
                           self, "e")

    @_handle_exceptions
563 564 565
    def new_graph_property(self, type, val=None):
        """Create a new graph property map of type `type`, and return it. If
        `val` is not None, the property is initialized to its value."""
566
        prop = PropertyMap(new_graph_property(_type_alias(type),
567 568
                                              self.__graph.GetGraphIndex()),
                           self, "g", lambda k: k.__graph)
569 570 571
        if val != None:
            prop[self] = val
        return prop
572 573 574 575

    # property map copying
    @_handle_exceptions
    @_require("src", PropertyMap)
576 577 578 579 580 581
    @_require("tgt", (PropertyMap, type(None)))
    def copy_property(self, src, tgt=None, g=None):
        """Copy contents of `src` property to `tgt` property. If `tgt` is None,
        then a new property map of the same type is created, and returned. The
        optional parameter g specifices the (identical) source graph to copy
        properties from (defaults to self).
582
        """
583 584 585 586 587 588
        if tgt == None:
            tgt = self.new_property(src.key_type(), src.value_type())
            ret = tgt
        else:
            ret = None

589 590 591 592 593 594 595 596 597
        if src.key_type() != tgt.key_type():
            raise GraphError("source and target properties must have the same" +
                             " key type")
        if g == None:
            g = self
        if g != self:
            g.stash_filter()
        self.stash_filter()
        if src.key_type() == "v":
598 599
            self.__graph.CopyVertexProperty(g.__graph, _prop("v", g, src),
                                            _prop("v", g, tgt))
600
        elif src.key_type() == "e":
601 602
            self.__graph.CopyEdgeProperty(g.__graph, _prop("e", g, src),
                                            _prop("e", g, tgt))
603 604 605 606 607
        else:
            tgt[self] = src[g]
        self.pop_filter()
        if g != self:
            g.pop_filter()
608
        return ret
609

610 611 612 613 614 615 616 617 618
    # degree property map
    @_handle_exceptions
    @_limit_args({"deg":["in","out","total"]})
    def degree_property_map(self, deg):
        """Create and return a vertex property map containing the degree type
        given by `deg`"""
        return PropertyMap(self.__graph.DegreeMap(deg), self, "v")


619 620
    # I/O operations
    # ==============
621 622 623

    @_handle_exceptions
    def load(self, filename, format="auto"):
624 625
        """Load graph from 'filename' (which can also be a file-like
        object). The format is guessed from the file name, or can be specified
626 627 628 629
        by 'format', which can be either 'xml' or 'dot'. Alternatively, filename
        can be file-like object."""
        if type(filename) == str:
            filename = os.path.expanduser(filename)
630 631 632 633 634 635 636 637
        if format == 'auto' and isinstance(filename, str):
            if filename.endswith(".xml") or filename.endswith(".xml.gz") or \
                   filename.endswith(".xml.bz2"):
                format = "xml"
            elif filename.endswith(".dot") or filename.endswith(".dot.gz") or \
                     filename.endswith(".dot.bz2"):
                format = "dot"
            else:
638
                libcore.raise_error\
639 640 641 642
                    ("cannot determine file format of: " + filename )
        elif format == "auto":
            format = "xml"
        if isinstance(filename, str):
643
            props = self.__graph.ReadFromFile(filename, None, format)
644
        else:
645 646 647 648 649 650 651 652
            props = self.__graph.ReadFromFile("", filename, format)
        for name, prop in props[0].iteritems():
            self.vertex_properties[name] = PropertyMap(prop, self, "v")
        for name, prop in props[1].iteritems():
            self.edge_properties[name] = PropertyMap(prop, self, "e")
        for name, prop in props[2].iteritems():
            self.graph_properties[name] = PropertyMap(prop, self, "g",
                                                      lambda k: k.__graph)
653 654 655 656

    @_handle_exceptions
    def save(self, filename, format="auto"):
        """Save graph to file. The format is guessed from the 'file' name, or
657 658 659 660
        can be specified by 'format', which can be either 'xml' or
        'dot'. Alternatively, filename can be file-like object."""
        if type(filename) == str:
            filename = os.path.expanduser(filename)
661 662 663 664 665 666 667 668
        if format == 'auto' and isinstance(filename, str):
            if filename.endswith(".xml") or filename.endswith(".xml.gz") or \
                   filename.endswith(".xml.bz2"):
                format = "xml"
            elif filename.endswith(".dot") or filename.endswith(".dot.gz") or \
                     filename.endswith(".dot.bz2"):
                format = "dot"
            else:
669
                libcore.raise_error\
670 671 672
                    ("cannot determine file format of: " + filename )
        elif format == "auto":
            format = "xml"
673 674
        props = [(name[1], prop._PropertyMap__map) for name,prop in \
                 self.__properties.iteritems()]
675
        if isinstance(filename, str):
676
            self.__graph.WriteToFile(filename, None, format, props)
677
        else:
678
            self.__graph.WriteToFile("", filename, format, props)
679

680 681
    # Directedness
    # ============
682 683 684 685 686 687 688 689 690 691 692

    @_handle_exceptions
    def set_directed(self, is_directed):
        """Set the directedness of the graph"""
        self.__graph.SetDirected(is_directed)

    @_handle_exceptions
    def is_directed(self):
        """Get the directedness of the graph"""
        return self.__graph.GetDirected()

693 694 695
    # Reversedness
    # ============

696 697 698 699 700 701 702 703 704 705 706
    @_handle_exceptions
    def set_reversed(self, is_reversed):
        """Reverse the direction of the edges, if 'reversed' is True, or
        maintain the original direction otherwise."""
        self.__graph.SetReversed(is_reversed)

    @_handle_exceptions
    def is_reversed(self):
        """Return 'True' if the edges are reversed, and 'False' otherwise."""
        return self.__graph.GetReversed()

707 708 709
    # Filtering
    # =========

710 711
    @_handle_exceptions
    def set_vertex_filter(self, property, inverted=False):
712 713 714 715 716
        """Choose vertex boolean filter property. Only the vertices with value
        different than zero are kept in the filtered graph. If the 'inverted'
        option is supplied with value 'True', only the vertices with value zero
        are kept. If the supplied property is 'None', any previous filtering is
        removed."""
717
        self.__graph.SetVertexFilterProperty(_prop("v", self, property), inverted)
718
        self.__filter_state["vertex_filter"] = (property, inverted)
719 720 721

    @_handle_exceptions
    def get_vertex_filter(self):
722 723
        """Get the vertex filter property and whether or not it is inverted"""
        return self.__filter_state["vertex_filter"]
724 725 726

    @_handle_exceptions
    def set_edge_filter(self, property, inverted=False):
727 728 729 730 731
        """Choose edge boolean filter property. Only the edges with value
        different than zero are kept in the filtered graph. If the 'inverted'
        option is supplied with value 'True', only the edges with value zero are
        kept. If the supplied property is 'None', any previous filtering is
        removed."""
732
        self.__graph.SetEdgeFilterProperty(_prop("e", self, property), inverted)
733
        self.__filter_state["edge_filter"] = (property, inverted)
734 735

    @_handle_exceptions
736
    def get_edge_filter(self):
737 738
        """Get the edge filter property and whether or not it is inverted"""
        return self.__filter_state["edge_filter"]
739 740 741 742 743 744

    @_handle_exceptions
    def purge_vertices(self):
        """Remove all vertices of the graph which are currently being filtered
        out, and return to the unfiltered state"""
        self.__graph.PurgeVertices()
745
        self.__graph.SetVertexFilterProperty(None)
746 747 748 749 750 751

    @_handle_exceptions
    def purge_edges(self):
        """Remove all edges of the graph which are currently being filtered out,
        and return to the unfiltered state"""
        self.__graph.PurgeEdges()
752
        self.__graph.SetEdgeFilterProperty(None)
753

754
    @_handle_exceptions
755 756
    def stash_filter(self, edge=False, vertex=False,
                     directed=False, reversed=False, all=True):
757 758
        """Stash current filter state and recover unfiltered graph. The optional
        keyword arguments specify which type of filter should be stashed."""
759 760
        if edge or vertex or directed or reversed:
            all = False
761
        self.__stashed_filter_state.append(self.__filter_state)
762
        if libcore.graph_filtering_enabled():
763 764 765 766 767 768 769 770 771 772
            if vertex or all:
                self.set_vertex_filter(None)
            if edge or all:
                self.set_edge_filter(None)
        if directed or all:
            self.set_directed(True)
        if reversed or all:
            self.set_reversed(False)

    @_handle_exceptions
773
    def pop_filter(self, edge=False, vertex=False,
774 775 776
                   directed=False, reversed=False, all=False):
        """Pop last stashed filter state. The optional keyword arguments specify
        which type of filter should be recovered."""
777 778
        if edge or vertex or directed or reversed:
            all = False
779 780 781 782 783 784 785 786 787 788 789 790
        state = self.__stashed_filter_state.pop()
        if libcore.graph_filtering_enabled():
            if vertex or all:
                self.set_vertex_filter(state["vertex_filter"][0],
                                       state["vertex_filter"][1])
            if edge or all:
                self.set_edge_filter(state["edge_filter"][0],
                                     state["edge_filter"][1])
        if directed or all:
            self.set_directed(state["directed"])
        if reversed or all:
            self.set_reversed(state["reversed"])
791 792

    @_handle_exceptions
793 794 795 796 797 798 799
    def get_filter_state(self):
        """Return a copy of the filter state of the graph."""
        return copy.copy(self.__filter_state)

    @_handle_exceptions
    def set_filter_state(self, state):
        """Set the filter state of the graph."""
800 801 802 803 804
        if libcore.graph_filtering_enabled():
            self.set_vertex_filter(state["vertex_filter"][0],
                                   state["vertex_filter"][1])
            self.set_edge_filter(state["edge_filter"][0],
                                 state["edge_filter"][1])
805 806 807
        self.set_directed(state["directed"])
        self.set_reversed(state["reversed"])

808
    # Basic graph statistics
809
    # ======================
810 811 812 813 814 815 816 817 818 819 820

    @_handle_exceptions
    def num_vertices(self):
        """Get the number of vertices."""
        return self.__graph.GetNumberOfVertices()

    @_handle_exceptions
    def num_edges(self):
        """Get the number of edges"""
        return self.__graph.GetNumberOfEdges()

821
    # Pickling support
822 823
    # ================

824 825 826
    def __getstate__(self):
        state = dict()
        sio = StringIO()
Tiago Peixoto's avatar
Tiago Peixoto committed
827 828 829 830 831 832 833 834 835
        if libcore.graph_filtering_enabled():
            if self.get_vertex_filter()[0] != None:
                self.vertex_properties["_Graph__pickle__vfilter"] = \
                    self.get_vertex_filter()[0]
                state["vfilter"] = self.get_vertex_filter()[1]
            if self.get_edge_filter()[0] != None:
                self.edge_properties["_Graph__pickle__efilter"] = \
                    self.get_edge_filter()[0]
                state["efilter"] = self.get_edge_filter()[1]
836 837 838
        self.save(sio, "xml")
        state["blob"] = sio.getvalue()
        return state
Tiago Peixoto's avatar
Tiago Peixoto committed
839

840 841 842 843 844 845
    def __setstate__(self, state):
        self.__init__()
        blob = state["blob"]
        if blob != "":
            sio = StringIO(blob)
            self.load(sio, "xml")
Tiago Peixoto's avatar
Tiago Peixoto committed
846 847 848 849 850 851
        if state.has_key("vfilt"):
            vprop = self.vertex_properties["_Graph__pickle__vfilter"]
            self.set_vertex_filter(vprop, state["vfilt"])
        if state.has_key("efilt"):
            eprop = self.edge_properties["_Graph__pickle__efilter"]
            self.set_edge_filter(vprop, state["efilt"])
852

853 854 855 856 857 858
def load_graph(filename, format="auto"):
    """Load a graph from file"""
    g = Graph()
    g.load(filename, format=format)
    return g

859 860
def value_types():
    """Return a list of possible properties value types"""
861
    return libcore.get_property_types()
862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894

# Vertex and Edge Types
# =====================
from libgraph_tool_core import Vertex, Edge

def _out_neighbours(self):
    """Return an iterator over the out-neighbours"""
    for e in self.out_edges():
        yield e.target()
Vertex.out_neighbours = _out_neighbours

def _in_neighbours(self):
    """Return an iterator over the in-neighbours"""
    for e in self.in_edges():
        yield e.source()
Vertex.in_neighbours = _in_neighbours

def _all_edges(self):
    """Return an iterator over all edges (both in or out)"""
    for e in self.out_edges():
        yield e
    for e in self.in_edges():
        yield e
Vertex.all_edges = _all_edges

def _all_neighbours(self):
    """Return an iterator over all neighbours (both in or out)"""
    for v in self.out_neighbours():
        yield v
    for v in self.in_neighbours():
        yield v
Vertex.all_neighbours = _all_neighbours

895 896 897 898 899
def _iter(self):
    """Iterate over the source and target"""
    for v in [self.source(), self.target()]:
        yield v
Edge.__iter__ = _iter