From 6dc7db88053e1bdb81276159b7e6d12485cdd782 Mon Sep 17 00:00:00 2001
From: Anshul Singhvi <anshulsinghvi@gmail.com>
Date: Thu, 25 Apr 2024 20:48:51 -0400
Subject: [PATCH] Allow Shapefile to read `.zip` files

We should subscribe to the FileIO interface so we can deal with streaming files as well...
---
 Project.toml               | 16 +++++++++-------
 ext/ShapefileZipFileExt.jl | 39 ++++++++++++++++++++++++++++++++++++++
 src/Shapefile.jl           | 17 +++++++++++++++++
 src/table.jl               |  5 +++++
 4 files changed, 70 insertions(+), 7 deletions(-)
 create mode 100644 ext/ShapefileZipFileExt.jl

diff --git a/Project.toml b/Project.toml
index ea45d57..2b78e69 100644
--- a/Project.toml
+++ b/Project.toml
@@ -14,6 +14,14 @@ OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
+[weakdeps]
+Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+ZipFile = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea"
+
+[extensions]
+ShapefileMakieExt = "Makie"
+ShapefileZipFileExt = "ZipFile"
+
 [compat]
 DBFTables = "1.2"
 Extents = "0.1"
@@ -21,18 +29,12 @@ GeoFormatTypes = "0.4"
 GeoInterface = "1.0"
 GeoInterfaceMakie = "0.1"
 GeoInterfaceRecipes = "1.0"
-Makie = "0.20"
+Makie = "0.20, 0.21"
 OrderedCollections = "1"
 RecipesBase = "1"
 Tables = "0.2, 1"
 julia = "1.9"
 
-[weakdeps]
-Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
-
-[extensions]
-ShapefileMakieExt = "Makie"
-
 [extras]
 ArchGDAL = "c9ce4bd3-c3d5-55b8-8973-c0e20141b8c3"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
diff --git a/ext/ShapefileZipFileExt.jl b/ext/ShapefileZipFileExt.jl
new file mode 100644
index 0000000..7a05053
--- /dev/null
+++ b/ext/ShapefileZipFileExt.jl
@@ -0,0 +1,39 @@
+module ShapefileZipFileExt
+import ZipFile, Shapefile
+import Shapefile: _read_shp_from_zipfile
+function _read_shp_from_zipfile(zipfile)
+  r = ZipFile.Reader(zipfile)
+  # need to get dbx
+  shpdata, shxdata, dbfdata, prjdata = nothing, nothing, nothing, nothing
+  for f in r.files
+    fn = f.name
+    lfn = lowercase(fn)
+    if endswith(lfn, ".shp")
+      shpdata = IOBuffer(read(f))
+    elseif endswith(lfn, ".shx")
+      shxdata = read(f, Shapefile.IndexHandle)
+    elseif endswith(lfn, ".dbf")
+      dbfdata = Shapefile.DBFTables.Table(IOBuffer(read(f)))
+    elseif endswith(lfn, "prj")
+      prjdata = try
+        Shapefile.GeoFormatTypes.ESRIWellKnownText(Shapefile.GeoFormatTypes.CRS(), read(f, String))
+      catch
+        @warn "Projection file $zipfile/$lfn appears to be corrupted. `nothing` used for `crs`"
+        nothing 
+      end
+    end
+  end
+  close(r)
+  @assert shpdata !== nothing
+  shp = if shxdata !== nothing # we have shxdata/index 
+    read(shpdata, Shapefile.Handle, shxdata)
+  else
+    read(shpdata, Shapefile.Handle)
+  end 
+  if prjdata !== nothing
+    shp.crs = prjdata 
+  end 
+  return Shapefile.Table(shp, dbfdata)
+end 
+
+end
\ No newline at end of file
diff --git a/src/Shapefile.jl b/src/Shapefile.jl
index dac11c7..97a0abf 100644
--- a/src/Shapefile.jl
+++ b/src/Shapefile.jl
@@ -70,4 +70,21 @@ include("extent.jl")
 include("plotrecipes.jl")
 include("writer.jl")
 
+function __init__()
+    # Register an error hint, so that if a user tries to read a zipfile and fails, they get a helpful error message
+    # that includes the ShapefileZipFileExt package.
+    Base.Experimental.register_error_hint(MethodError) do io, exc, argtypes, kwargs
+        if exc.f == _read_shp_from_zipfile
+            if isnothing(Base.get_extension(Shapefile, :ShapefileZipFileExt))
+                print(io, "\nPlease load the ")
+                printstyled(io, "ZipFile", color=:cyan)
+                println(io, " package to read zipfiles into Shapefile.Table objects.")
+                println(io, "You can do this by typing: ")
+                printstyled(io, "using ZipFile", color=:cyan, bold = true)
+                println(io, "\ninto your REPL or code.")
+            end
+        end
+    end
+end
+
 end # module
diff --git a/src/table.jl b/src/table.jl
index fbb8b50..fd9204d 100644
--- a/src/table.jl
+++ b/src/table.jl
@@ -67,6 +67,9 @@ function Table(shp::Handle{T}, dbf::DBFTables.Table) where {T}
     Table{T}(shp, dbf)
 end
 function Table(path::AbstractString)
+    if endswith(path, ".zip")
+        return _read_shp_from_zipfile(path)
+    end
     paths = _shape_paths(path)
     isfile(paths.shp) || throw(ArgumentError("File not found: $(paths.dbf)"))
     isfile(paths.dbf) || throw(ArgumentError("File not found: $(paths.dbf)"))
@@ -80,6 +83,8 @@ function Table(path::AbstractString)
     return Shapefile.Table(shp, dbf)
 end
 
+function _read_shp_from_zipfile end
+
 getshp(t::Table) = getfield(t, :shp)
 getdbf(t::Table) = getfield(t, :dbf)