From 9e5fa98d27dfa4d7e0227bf2f1995e247d57f4fd Mon Sep 17 00:00:00 2001
From: Joris Van den Bossche <jorisvandenbossche@gmail.com>
Date: Thu, 27 Sep 2018 09:28:49 +0200
Subject: [PATCH] BUG/TST: fix align to preserve metadata + more tests (#829)

---
 geopandas/geoseries.py               |  6 +--
 geopandas/tests/test_geodataframe.py | 60 +++++++++++++++++++++++++++-
 geopandas/tests/test_geoseries.py    | 23 +++++++++++
 3 files changed, 82 insertions(+), 7 deletions(-)

diff --git a/geopandas/geoseries.py b/geopandas/geoseries.py
index d90d90c..c2032d7 100644
--- a/geopandas/geoseries.py
+++ b/geopandas/geoseries.py
@@ -238,11 +238,7 @@ class GeoSeries(GeoPandasBase, Series):
                                                    level=level, copy=copy,
                                                    fill_value=fill_value,
                                                    **kwargs)
-        if isinstance(other, GeoSeries):
-            return GeoSeries(left), GeoSeries(right)
-        else: # It is probably a Series, let's keep it that way
-            return GeoSeries(left), right
-
+        return left, right
 
     def __contains__(self, other):
         """Allow tests of the form "geom in s"
diff --git a/geopandas/tests/test_geodataframe.py b/geopandas/tests/test_geodataframe.py
index a41b758..8cccd9c 100644
--- a/geopandas/tests/test_geodataframe.py
+++ b/geopandas/tests/test_geodataframe.py
@@ -17,8 +17,8 @@ import pytest
 from pandas.util.testing import (
     assert_frame_equal, assert_index_equal, assert_series_equal)
 from geopandas.tests.util import (
-    assert_geoseries_equal, connect, create_postgis, PACKAGE_DIR,
-    validate_boro_df)
+    connect, create_postgis, PACKAGE_DIR, validate_boro_df)
+from geopandas.testing import assert_geodataframe_equal, assert_geoseries_equal
 
 
 class TestDataFrame:
@@ -210,6 +210,62 @@ class TestDataFrame:
             assert i == r['geometry'].x
             assert i == r['geometry'].y
 
+    def test_align(self):
+        df = self.df2
+
+        res1, res2 = df.align(df)
+        assert_geodataframe_equal(res1, df)
+        assert_geodataframe_equal(res2, df)
+
+        res1, res2 = df.align(df.copy())
+        assert_geodataframe_equal(res1, df)
+        assert_geodataframe_equal(res2, df)
+
+        # assert crs is / is not preserved on mixed dataframes
+        df_nocrs = df.copy()
+        df_nocrs.crs = None
+        res1, res2 = df.align(df_nocrs)
+        assert_geodataframe_equal(res1, df)
+        assert res1.crs is not None
+        assert_geodataframe_equal(res2, df_nocrs)
+        assert res2.crs is None
+
+        # mixed GeoDataFrame / DataFrame
+        df_nogeom = pd.DataFrame(df.drop('geometry', axis=1))
+        res1, res2 = df.align(df_nogeom, axis=0)
+        assert_geodataframe_equal(res1, df)
+        assert type(res2) == pd.DataFrame
+        assert_frame_equal(res2, df_nogeom)
+
+        # same as above but now with actual alignment
+        df1 = df.iloc[1:].copy()
+        df2 = df.iloc[:-1].copy()
+
+        exp1 = df.copy()
+        exp1.iloc[0] = np.nan
+        exp2 = df.copy()
+        exp2.iloc[-1] = np.nan
+        res1, res2 = df1.align(df2)
+        assert_geodataframe_equal(res1, exp1)
+        assert_geodataframe_equal(res2, exp2)
+
+        df2_nocrs = df2.copy()
+        df2_nocrs.crs = None
+        exp2_nocrs = exp2.copy()
+        exp2_nocrs.crs = None
+        res1, res2 = df1.align(df2_nocrs)
+        assert_geodataframe_equal(res1, exp1)
+        assert res1.crs is not None
+        assert_geodataframe_equal(res2, exp2_nocrs)
+        assert res2.crs is None
+
+        df2_nogeom = pd.DataFrame(df2.drop('geometry', axis=1))
+        exp2_nogeom = pd.DataFrame(exp2.drop('geometry', axis=1))
+        res1, res2 = df1.align(df2_nogeom, axis=0)
+        assert_geodataframe_equal(res1, exp1)
+        assert type(res2) == pd.DataFrame
+        assert_frame_equal(res2, exp2_nogeom)
+
     def test_to_json(self):
         text = self.df.to_json()
         data = json.loads(text)
diff --git a/geopandas/tests/test_geoseries.py b/geopandas/tests/test_geoseries.py
index ec96775..07389be 100644
--- a/geopandas/tests/test_geoseries.py
+++ b/geopandas/tests/test_geoseries.py
@@ -102,6 +102,29 @@ class TestSeries:
         assert a1['B'].equals(a2['B'])
         assert a1['C'].is_empty
 
+    def test_align_crs(self):
+        a1 = self.a1
+        a1.crs = {'init': 'epsg:4326', 'no_defs': True}
+        a2 = self.a2
+        a2.crs = {'init': 'epsg:31370', 'no_defs': True}
+
+        res1, res2 = a1.align(a2)
+        assert res1.crs == {'init': 'epsg:4326', 'no_defs': True}
+        assert res2.crs == {'init': 'epsg:31370', 'no_defs': True}
+
+        a2.crs = None
+        res1, res2 = a1.align(a2)
+        assert res1.crs == {'init': 'epsg:4326', 'no_defs': True}
+        assert res2.crs is None
+
+    def test_align_mixed(self):
+        a1 = self.a1
+        s2 = pd.Series([1, 2], index=['B', 'C'])
+        res1, res2 = a1.align(s2)
+
+        exp2 = pd.Series([BaseGeometry(), 1, 2], dtype=object, index=['A', 'B', 'C'])
+        assert_series_equal(res2, exp2)
+
     def test_geom_almost_equals(self):
         # TODO: test decimal parameter
         assert np.all(self.g1.geom_almost_equals(self.g1))