Attachments you submit will be routed for moderation. If you have an account, please log in first.

Ticket #468: nx_shp.py

File nx_shp.py, 2.6 KB (added by benwreilly@…, 3 years ago)

shapefile read support

Line 
1"""
2*********
3Shapefile
4*********
5
6Read shapefiles to generate Networkx DiGraphs.
7
8"The Esri Shapefile or simply a shapefile is a popular geospatial vector
9data format for geographic information systems software. It is developed
10and regulated by Esri as a (mostly) open specification for data
11interoperability among Esri and other software products."
12See http://en.wikipedia.org/wiki/Shapefile for additional information.
13
14"""
15__author__ = """Ben Reilly (benwreilly@gmail.com)"""
16#    Copyright (C) 2004-2010 by
17#    Ben Reilly <benwreilly@gmail.com>
18#    Aric Hagberg <hagberg@lanl.gov>
19#    Dan Schult <dschult@colgate.edu>
20#    Pieter Swart <swart@lanl.gov>
21#    All rights reserved.
22#    BSD license.
23
24__all__ = ['read_shp']
25
26import networkx as nx
27
28def read_shp(path):
29    """Read graph in SHP format from path.
30
31    "The Esri Shapefile or simply a shapefile is a popular geospatial vector
32    data format for geographic information systems software [1]_."
33
34    Parameters
35    ----------
36    path : file or string
37       File or filename to read.
38
39    Returns
40    -------
41    G : NetworkX graph
42
43    Examples
44    --------
45    >>> G=nx.read_shp('test.shp')
46   
47    References
48    ----------
49    .. [1] http://en.wikipedia.org/wiki/Shapefile
50    """
51    try:
52        from osgeo import ogr
53    except ImportError:
54        raise ImportError("read_shp() requires OGR: http://www.gdal.org/")
55   
56    net = nx.DiGraph()
57   
58    def getfieldinfo(lyr, feature, flds):
59            f = feature
60            return [f.GetField(f.GetFieldIndex(x)) for x in flds]
61           
62    def addlyr(lyr, fields):
63        for findex in xrange(lyr.GetFeatureCount()):
64            f = lyr.GetFeature(findex)
65            flddata = getfieldinfo(lyr, f, fields)
66            g = f.geometry()
67            attributes = dict(zip(fields, flddata))
68            attributes["ShpName"] = lyr.GetName()
69            if g.GetGeometryType() == 1: #point
70                net.add_node((g.GetPoint_2D(0)), attributes)
71            if g.GetGeometryType() == 2: #linestring
72                last = g.GetPointCount() - 1
73                net.add_edge(g.GetPoint_2D(0), g.GetPoint_2D(last), attributes)
74               
75    if isinstance(path, str):
76        shp = ogr.Open(path)
77        lyrcount = shp.GetLayerCount()
78        for lyrindex in xrange(lyrcount):
79            lyr = shp.GetLayerByIndex(lyrindex)
80            flds = [x.GetName() for x in lyr.schema]
81            addlyr(lyr, flds)
82   
83    return net
84   
85# fixture for nose tests
86def setup_module(module):
87    from nose import SkipTest
88    try:
89        import ogr
90    except:
91        raise SkipTest("OGR not available")