-
Notifications
You must be signed in to change notification settings - Fork 24
/
geodataframe.py
129 lines (106 loc) · 4.19 KB
/
geodataframe.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
import pandas as pd
from ._optional_imports import gp
from .geometry import GeometryDtype
from .geoseries import GeoSeries, _MaybeGeoSeries
class _MaybeGeoDataFrame(pd.DataFrame):
def __new__(cls, *args, **kwargs):
try:
return GeoDataFrame(*args, **kwargs)
except ValueError:
# No geometry compatible columns
return pd.DataFrame(*args, **kwargs)
class GeoDataFrame(pd.DataFrame):
# properties to propagate
_metadata = ['_geometry']
# In Pandas 2.1 will raise AttributeError.
# AttributeError: 'GeoDataFrame' object has no attribute '_geometry'
# Ref: https://github.com/pandas-dev/pandas/issues/51280
_geometry = None
def __init__(self, data=None, index=None, geometry=None, **kwargs):
# Call pandas constructor, always copy
kwargs.pop("copy", None)
super().__init__(data, index=index, copy=True, **kwargs)
# Replace pd.Series of GeometryArrays with GeoSeries.
first_geometry_col = None
for col in self.columns:
if (isinstance(self[col].dtype, GeometryDtype) or
gp and isinstance(self[col].dtype, gp.array.GeometryDtype)):
self[col] = GeoSeries(self[col])
first_geometry_col = first_geometry_col or col
if first_geometry_col is None:
raise ValueError(
"A spatialpandas GeoDataFrame must contain at least one spatialpandas "
"GeometryArray column"
)
if geometry is None:
if isinstance(data, GeoDataFrame) and data._has_valid_geometry():
geometry = data._geometry
elif gp and isinstance(data, gp.GeoDataFrame):
try:
geometry = data.geometry.name
except AttributeError:
# Geometry column not set
pass
if geometry is None:
geometry = first_geometry_col
self._geometry = None
if geometry is not None:
self.set_geometry(geometry, inplace=True)
@property
def _constructor(self):
return _MaybeGeoDataFrame
@property
def _constructor_sliced(self):
return _MaybeGeoSeries
def set_geometry(self, geometry, inplace=False):
if (geometry not in self or
not isinstance(self[geometry].dtype, GeometryDtype)):
raise ValueError(
"The geometry argument must be the name of a spatialpandas "
"geometry column in the spatialpandas GeoDataFrame"
)
if inplace:
self._geometry = geometry
return self
else:
return GeoDataFrame(self, geometry=geometry)
def _has_valid_geometry(self):
if (self._geometry is not None and
self._geometry in self and
isinstance(self[self._geometry].dtype, GeometryDtype)):
return True
else:
return False
@property
def geometry(self):
if not self._has_valid_geometry():
raise ValueError(
"GeoDataFrame has no active geometry column.\n"
"The active geometry column should be set using the set_geometry "
"method."
)
return self[self._geometry]
def to_geopandas(self):
from geopandas import GeoDataFrame as gp_GeoDataFrame
gdf = gp_GeoDataFrame(
{col: s.array.to_geopandas() if isinstance(s.dtype, GeometryDtype) else s
for col, s in self.items()},
index=self.index,
)
if self._has_valid_geometry():
gdf.set_geometry(self._geometry, inplace=True)
return gdf
@property
def cx(self):
from .geometry.base import _CoordinateIndexer
return _CoordinateIndexer(
self.geometry.array, parent=self
)
def build_sindex(self, **kwargs):
self.geometry.build_sindex(**kwargs)
return self
def _ensure_type(self, obj):
# Override because a GeoDataFrame operation may result in a regular DataFrame,
# and that's ok
assert isinstance(obj, type(self)) or isinstance(self, type(obj)), type(obj)
return obj