Add WKT/PROJ CRS support via pyproj

CRS can now be specified as WKT strings, PROJ strings, or EPSG
integers. pyproj (lazy import) resolves between them:

Read side:
- crs_wkt attr is populated by resolving EPSG -> WKT via pyproj
- Falls back gracefully if pyproj is not installed (EPSG still works)

Write side:
- crs= parameter on write_geotiff accepts int (EPSG), WKT string,
  or PROJ string. String inputs are resolved to EPSG via
  pyproj.CRS.from_user_input().to_epsg().
- DataArray with crs_wkt attr (no integer crs) is also handled:
  the WKT is resolved to EPSG for the GeoKeyDirectory.

This means files with user-defined CRS no longer lose their spatial
reference when round-tripped, as long as pyproj can resolve the
WKT/PROJ string to an EPSG code.

5 new tests: WKT from EPSG, write with WKT string, write with PROJ
string, crs_wkt attr round-trip, and no-CRS baseline.

Author: brendancol

xrspatial/geotiff/__init__.py

@@ -23,6 +23,20 @@
 __all__ = ['read_geotiff', 'write_geotiff', 'open_cog', 'read_geotiff_dask']
 
 
+def _wkt_to_epsg(wkt_or_proj: str) -> int | None:
+    """Try to extract an EPSG code from a WKT or PROJ string.
+
+    Returns None if pyproj is not installed or the string can't be parsed.
+    """
+    try:
+        from pyproj import CRS
+        crs = CRS.from_user_input(wkt_or_proj)
+        epsg = crs.to_epsg()
+        return epsg
+    except Exception:
+        return None
+
+
 def _geo_to_coords(geo_info, height: int, width: int) -> dict:
     """Build y/x coordinate arrays from GeoInfo.
 
@@ -132,6 +146,8 @@ def read_geotiff(source: str, *, window=None,
     attrs = {}
     if geo_info.crs_epsg is not None:
         attrs['crs'] = geo_info.crs_epsg
+    if geo_info.crs_wkt is not None:
+        attrs['crs_wkt'] = geo_info.crs_wkt
     if geo_info.raster_type == RASTER_PIXEL_IS_POINT:
         attrs['raster_type'] = 'point'
 
@@ -214,7 +230,7 @@ def read_geotiff(source: str, *, window=None,
 
 
 def write_geotiff(data: xr.DataArray | np.ndarray, path: str, *,
-                  crs: int | None = None,
+                  crs: int | str | None = None,
                   nodata=None,
                   compression: str = 'deflate',
                   tiled: bool = True,
@@ -231,8 +247,10 @@ def write_geotiff(data: xr.DataArray | np.ndarray, path: str, *,
         2D raster data.
     path : str
         Output file path.
-    crs : int or None
-        EPSG code. If None and data is a DataArray, tries to read from attrs.
+    crs : int, str, or None
+        EPSG code (int), WKT string, or PROJ string. If None and data
+        is a DataArray, tries to read from attrs ('crs' for EPSG,
+        'crs_wkt' for WKT).
     nodata : float, int, or None
         NoData value.
     compression : str
@@ -252,18 +270,29 @@ def write_geotiff(data: xr.DataArray | np.ndarray, path: str, *,
         'min', 'max', 'median', 'mode', or 'cubic'.
     """
     geo_transform = None
-    epsg = crs
+    epsg = None
     raster_type = RASTER_PIXEL_IS_AREA
     x_res = None
     y_res = None
     res_unit = None
 
+    # Resolve crs argument: can be int (EPSG) or str (WKT/PROJ)
+    if isinstance(crs, int):
+        epsg = crs
+    elif isinstance(crs, str):
+        epsg = _wkt_to_epsg(crs)  # try to extract EPSG from WKT/PROJ
+
     if isinstance(data, xr.DataArray):
         arr = data.values
         if geo_transform is None:
             geo_transform = _coords_to_transform(data)
-        if epsg is None:
+        if epsg is None and crs is None:
             epsg = data.attrs.get('crs')
+            if epsg is None:
+                # Try resolving EPSG from a WKT string in attrs
+                wkt = data.attrs.get('crs_wkt')
+                if isinstance(wkt, str):
+                    epsg = _wkt_to_epsg(wkt)
         if nodata is None:
             nodata = data.attrs.get('nodata')
         if data.attrs.get('raster_type') == 'point':

xrspatial/geotiff/_geotags.py

@@ -107,10 +107,24 @@ class GeoInfo:
     vertical_datum: int | None = None
     vertical_units: str | None = None
     vertical_units_code: int | None = None
+    # WKT CRS string (resolved from EPSG via pyproj, or provided by caller)
+    crs_wkt: str | None = None
     # Raw geokeys dict for anything else
     geokeys: dict[int, int | float | str] = field(default_factory=dict)
 
 
+def _epsg_to_wkt(epsg: int) -> str | None:
+    """Resolve an EPSG code to a WKT string using pyproj.
+
+    Returns None if pyproj is not installed or the code is unknown.
+    """
+    try:
+        from pyproj import CRS
+        return CRS.from_epsg(epsg).to_wkt()
+    except Exception:
+        return None
+
+
 def _parse_geokeys(ifd: IFD, data: bytes | memoryview,
                    byte_order: str) -> dict[int, int | float | str]:
     """Parse the GeoKeyDirectory and resolve values from param tags.
@@ -385,6 +399,11 @@ def extract_geo_info(ifd: IFD, data: bytes | memoryview,
                     b = raw_cmap[2 * n_colors + i] / 65535.0
                     colormap.append((r, g, b, 1.0))
 
+    # Resolve EPSG -> WKT via pyproj if available
+    crs_wkt = None
+    if epsg is not None:
+        crs_wkt = _epsg_to_wkt(epsg)
+
     return GeoInfo(
         transform=transform,
         crs_epsg=epsg,
@@ -410,6 +429,7 @@ def extract_geo_info(ifd: IFD, data: bytes | memoryview,
         vertical_datum=vert_datum,
         vertical_units=vert_units_name,
         vertical_units_code=vert_units_code,
+        crs_wkt=crs_wkt,
         geokeys=geokeys,
     )
 

xrspatial/geotiff/tests/test_features.py

@@ -337,6 +337,77 @@ def test_angular_unit_lookup(self):
         assert LINEAR_UNITS[9002] == 'foot'
         assert LINEAR_UNITS[9003] == 'us_survey_foot'
 
+    def test_crs_wkt_from_epsg(self, tmp_path):
+        """crs_wkt is resolved from EPSG via pyproj."""
+        from xrspatial.geotiff._geotags import GeoTransform
+        arr = np.ones((4, 4), dtype=np.float32)
+        gt = GeoTransform(-120.0, 45.0, 0.001, -0.001)
+        path = str(tmp_path / 'wkt.tif')
+        write(arr, path, compression='none', tiled=False,
+              geo_transform=gt, crs_epsg=4326)
+
+        da = read_geotiff(path)
+        assert 'crs_wkt' in da.attrs
+        wkt = da.attrs['crs_wkt']
+        assert 'WGS 84' in wkt or '4326' in wkt
+
+    def test_write_with_wkt_string(self, tmp_path):
+        """crs= accepts a WKT string and resolves to EPSG."""
+        arr = np.ones((4, 4), dtype=np.float32)
+        wkt = ('GEOGCRS["WGS 84",DATUM["World Geodetic System 1984",'
+               'ELLIPSOID["WGS 84",6378137,298.257223563]],'
+               'CS[ellipsoidal,2],'
+               'AXIS["geodetic latitude (Lat)",north],'
+               'AXIS["geodetic longitude (Lon)",east],'
+               'UNIT["degree",0.0174532925199433],'
+               'ID["EPSG",4326]]')
+        path = str(tmp_path / 'wkt_in.tif')
+        write_geotiff(arr, path, crs=wkt, compression='none')
+
+        da = read_geotiff(path)
+        assert da.attrs['crs'] == 4326
+
+    def test_write_with_proj_string(self, tmp_path):
+        """crs= accepts a PROJ string."""
+        arr = np.ones((4, 4), dtype=np.float32)
+        path = str(tmp_path / 'proj_in.tif')
+        write_geotiff(arr, path, crs='+proj=utm +zone=18 +datum=NAD83',
+                      compression='none')
+
+        da = read_geotiff(path)
+        # pyproj should resolve this to EPSG:26918
+        assert da.attrs.get('crs') is not None
+
+    def test_crs_wkt_attr_round_trip(self, tmp_path):
+        """DataArray with crs_wkt attr (no int crs) round-trips."""
+        wkt = ('GEOGCRS["WGS 84",DATUM["World Geodetic System 1984",'
+               'ELLIPSOID["WGS 84",6378137,298.257223563]],'
+               'CS[ellipsoidal,2],'
+               'AXIS["geodetic latitude (Lat)",north],'
+               'AXIS["geodetic longitude (Lon)",east],'
+               'UNIT["degree",0.0174532925199433],'
+               'ID["EPSG",4326]]')
+        y = np.linspace(45.0, 44.0, 4)
+        x = np.linspace(-120.0, -119.0, 4)
+        da = xr.DataArray(np.ones((4, 4), dtype=np.float32),
+                          dims=['y', 'x'], coords={'y': y, 'x': x},
+                          attrs={'crs_wkt': wkt})
+        path = str(tmp_path / 'wkt_rt.tif')
+        write_geotiff(da, path, compression='none')
+
+        result = read_geotiff(path)
+        assert result.attrs['crs'] == 4326
+        assert 'crs_wkt' in result.attrs
+
+    def test_no_crs_no_wkt(self, tmp_path):
+        """File without CRS has no crs_wkt attr."""
+        arr = np.ones((4, 4), dtype=np.float32)
+        path = str(tmp_path / 'no_wkt.tif')
+        write(arr, path, compression='none', tiled=False)
+
+        da = read_geotiff(path)
+        assert 'crs_wkt' not in da.attrs
+
 
 # -----------------------------------------------------------------------
 # Resolution / DPI tags