Support reading encodings from .cpg files (all the various ways)

Author: JamesParrott

src/shapefile.py

@@ -3599,24 +3599,23 @@ def __init__(
         shapefile_path: str | PathLike[Any] = "",
         /,
         *,
-        encoding: str = "utf-8",  # | None = None,
+        encoding: str | None = None,
         encodingErrors: str = "strict",
         shp: _NoShpSentinel | BinaryFileT | None = _NoShpSentinel(),
         shx: BinaryFileT | None = None,
         dbf: BinaryFileT | None = None,
-        # cpg: BinaryFileT | None = None,
+        cpg: BinaryFileT | None = None,
         # Keep kwargs even though unused, to preserve PyShp 2.4 API
         **kwargs: Any,
     ):
         super().__init__()
         # Store encoding info to use if lazy loading DbfReader later.
-        # encoding_from_cpg =
-        self.encoding = encoding  # encoding_from_cpg if encoding is None else encoding
+        self._user_specified_encoding = encoding
         self.encodingErrors = encodingErrors
         self._shp = None
         self._shx = None
         self._dbf = None
-        # self._cpg = None
+        self._cpg = None
         self.shapeName = "Not specified"
         self.numShapes: int = 0
         self.path: str | os.PathLike[Any] | None = None
@@ -3635,15 +3634,15 @@ def __init__(
                 discarded_kwargs["shx"] = shx
             if dbf is not None:
                 discarded_kwargs["dbf"] = dbf
-            # if cpg is not None:
-            #     discarded_kwargs["cpg"] = cpg
+            if cpg is not None:
+                discarded_kwargs["cpg"] = cpg
             if discarded_kwargs:
                 raise TypeError(
                     "Please be specific about the shapefile you want to load. "
                     f"Got: {shapefile_path}, plus the following unusable "
                     " kwargs: {discarded_kwargs} which previous versions of PyShp ignored. \n"
                     "Only either: i) exactly one positional arg \n"
-                    "         or: ii) one or both of shp and dbf, optionally plus shx \n"  # and cpg kwargs\n"
+                    "         or: ii) one or both of shp and dbf, optionally plus shx and cpg kwargs\n"
                     "is currently supported.  All other kwargs may be set (or not). "
                 )
             self.path = shapefile_path
@@ -3670,6 +3669,7 @@ def __init__(
                         zipfileobj = self._download_binary_file_from_url(
                             url_info,
                             ".zip",
+                            add_to_exit_stack=False,
                             suppress_http_errors=False,
                         )
                         if zipfileobj is None:
@@ -3708,12 +3708,14 @@ def __init__(
             self._shp = self._seek_0_on_file_obj_wrap_or_open_from_name(".shp", shp)
             self._shx = self._seek_0_on_file_obj_wrap_or_open_from_name(".shx", shx)
 
+        self._cpg = self._seek_0_on_file_obj_wrap_or_open_from_name(".cpg", cpg)
         self._dbf = self._seek_0_on_file_obj_wrap_or_open_from_name(".dbf", dbf)
 
         # Load the files
         if self._shp:
             self._get_shp_reader()
         if self._dbf:
+            # Sets self.encoding
             self._get_dbf_reader()
         if self._shx:
             self._get_shx_reader()
@@ -3747,12 +3749,37 @@ def _get_dbf_reader(self) -> DbfReader:
             raise ShapefileException(
                 "DbfReader requires a .dbf file or file-like object."
             )
+        self._set_encoding()
         return DbfReader(
             dbf=self._dbf,
             encoding=self.encoding,
             encodingErrors=self.encodingErrors,
         )
 
+    @functools.cache
+    def _set_encoding(self) -> None:
+        if self._cpg is None:
+            encoding_from_cpg = ""
+        else:
+            encoding_from_cpg = (
+                self._cpg.read().decode().lower().replace("-", "_").strip()
+            )
+            if not encoding_from_cpg:
+                warnings.warn("Empty .cpg file (no encoding found). ")
+
+        if self._user_specified_encoding is None:
+            encoding = encoding_from_cpg
+        else:
+            encoding = self._user_specified_encoding.lower().replace("-", "_").strip()
+
+            if encoding_from_cpg and encoding != encoding_from_cpg:
+                warnings.warn(
+                    f"Specified encoding: {encoding} "
+                    "different to encoding read from "
+                    f".cpg file: {encoding_from_cpg}",
+                )
+        self.encoding = encoding or "utf-8"
+
     @property
     def shp_reader(self) -> ShpReader:
         return self._get_shp_reader()
@@ -3843,7 +3870,7 @@ def iterRecords(
     def _try_get_open_constituent_file(
         self,
         file: Path,
-        ext: Literal[".shp", ".shx", ".dbf"],
+        ext: Literal[".shp", ".shx", ".dbf", ".cpg"],
     ) -> IO[bytes] | None:
         """
         Attempts to open a .shp, .dbf or .shx file,
@@ -3863,15 +3890,18 @@ def _try_get_open_constituent_file(
 
     def _seek_0_on_file_obj_wrap_or_open_from_name(
         self,
-        ext: Literal[".shp", ".shx", ".dbf"],
+        ext: Literal[".shp", ".shx", ".dbf", ".cpg"],
         file: BinaryFileT | None,
     ) -> None | IO[bytes]:
         if file is None:
             return None
 
         if isinstance(file, (str, PathLike)):
+            # Added to exit stack if opened.
             return self._try_get_open_constituent_file(Path(file), ext)
 
+        # Other user-opened file objects not added to exit stack.
+        # The user must close them.
         if hasattr(file, "read"):
             # Copy if required
             try:
@@ -3888,7 +3918,8 @@ def _seek_0_on_file_obj_wrap_or_open_from_name(
     def _download_binary_file_from_url(
         self,
         urlinfo: SplitResult,
-        ext: Literal[".shp", ".shx", ".dbf", ".zip"],
+        ext: Literal[".shp", ".shx", ".dbf", ".cpg", ".zip"],
+        add_to_exit_stack: bool = True,
         suppress_http_errors: bool = True,
     ) -> tempfile._TemporaryFileWrapper[bytes] | None:
         sniffed_bytes, resp = _try_to_download_binary_file(
@@ -3900,26 +3931,19 @@ def _download_binary_file_from_url(
             return None
         # Use tempfile as source for url data.
         fileobj = _save_to_named_tmp_file(resp, initial_bytes=sniffed_bytes)
+        if add_to_exit_stack:
+            self.exit_stack.enter_context(fileobj)
         return fileobj
 
     def _load_from_url(self, urlinfo: SplitResult) -> None:
         # Shapefile is from a url
         # Download each file to temporary path and treat as normal shapefile path
         self._shp = self._download_binary_file_from_url(urlinfo, ".shp")
         self._shx = self._download_binary_file_from_url(urlinfo, ".shx")
+        self._cpg = self._download_binary_file_from_url(urlinfo, ".cpg")
         self._dbf = self._download_binary_file_from_url(urlinfo, ".dbf")
 
-        shp_or_dbf_loaded = False
-        if self._shx is not None:
-            self.exit_stack.enter_context(self._shx)
-        if self._shp is not None:
-            self.exit_stack.enter_context(self._shp)
-            shp_or_dbf_loaded = True
-        if self._dbf is not None:
-            self.exit_stack.enter_context(self._dbf)
-            shp_or_dbf_loaded = True
-
-        if not shp_or_dbf_loaded:
+        if self._shp is None and self._dbf is None:
             raise ShapefileException(
                 f"Failed to download .shp or .dbf from: {urlunsplit(urlinfo)}"
             )
@@ -3928,7 +3952,7 @@ def _load_file_from_zip_to_tmp_file(
         self,
         archive: zipfile.ZipFile,
         file: Path,
-        ext: Literal[".shp", ".shx", ".dbf"],
+        ext: Literal[".shp", ".shx", ".dbf", ".cpg"],
     ) -> tempfile._TemporaryFileWrapper[bytes] | None:
         for cased_ext in {ext.lower(), ext.upper(), ext}:
             try:
@@ -3982,7 +4006,7 @@ def _load_from_zipfileobj(
                 constituent_files = (
                     Path(name)
                     for name in archive.namelist()
-                    if name.lower().endswith((".shp", ".dbf", ".shx"))
+                    if name.lower().endswith((".shp", ".dbf", ".shx", ".cpg"))
                 )
 
                 def without_ext(path: Path) -> Path:
@@ -4005,6 +4029,7 @@ def without_ext(path: Path) -> Path:
             # Try to extract file-like objects from zipfile
             self._shp = self._load_file_from_zip_to_tmp_file(archive, shapefile, ".shp")
             self._shx = self._load_file_from_zip_to_tmp_file(archive, shapefile, ".shx")
+            self._cpg = self._load_file_from_zip_to_tmp_file(archive, shapefile, ".cpg")
             self._dbf = self._load_file_from_zip_to_tmp_file(archive, shapefile, ".dbf")
 
     def load(self, file: str | os.PathLike[Any]) -> None:
@@ -4046,6 +4071,7 @@ def load_dbf(self, file: Path) -> None:
         """
         self._dbf = self._try_get_open_constituent_file(file, ".dbf")
         if self._dbf is not None:
+            self._cpg = self._try_get_open_constituent_file(file, ".cpg")
             self._get_dbf_reader()
 
     def __len__(self) -> int:

tests/test_shapefile.py

@@ -2194,8 +2194,7 @@ def test_read_latin1_shapefile(encoding, context, tmp_latin1_shapefile_shp):
     if encoding == "latin1":
         assert rec == [2, u'Ñandú']
 
-@pytest.mark.xfail(reason="Support for reading encodings from .cpg files not implemented yet")
-@pytest.mark.parametrize("encoding, context", ENCODINGS_AND_CONTEXTS[:1])
+@pytest.mark.parametrize("encoding, context", ENCODINGS_AND_CONTEXTS)
 def test_read_latin1_shapefile_cpg_file(encoding, context, tmp_latin1_shapefile_shp):
     """ Extend the smoke test in README.md doctests """