From d9f917a59a922c6981302f99fad8d94013bcf5cd Mon Sep 17 00:00:00 2001 From: Mike Taves Date: Wed, 16 Feb 2022 00:00:58 +1300 Subject: [PATCH] Support path-like objects for reading/writing --- shapefile.py | 29 ++++++++++++++++++++++++----- test_shapefile.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 5 deletions(-) diff --git a/shapefile.py b/shapefile.py index 2132fd5..133a20f 100644 --- a/shapefile.py +++ b/shapefile.py @@ -161,6 +161,24 @@ def u(v, encoding='utf-8', encodingErrors='strict'): def is_string(v): return isinstance(v, basestring) +if sys.version_info[0:2] >= (3, 6): + def pathlike_obj(path): + if isinstance(path, os.PathLike): + return os.fsdecode(path) + else: + return path +else: + def pathlike_obj(path): + if is_string(path): + return path + elif hasattr(path, "__fspath__"): + return path.__fspath__() + else: + try: + return str(path) + except: + return path + # Begin @@ -930,14 +948,14 @@ def __init__(self, *args, **kwargs): self.encodingErrors = kwargs.pop('encodingErrors', 'strict') # See if a shapefile name was passed as the first argument if len(args) > 0: - if is_string(args[0]): - path = args[0] - + path = pathlike_obj(args[0]) + if is_string(path): + if '.zip' in path: # Shapefile is inside a zipfile if path.count('.zip') > 1: # Multiple nested zipfiles - raise ShapefileException('Reading from multiple nested zipfiles is not supported: %s' % args[0]) + raise ShapefileException('Reading from multiple nested zipfiles is not supported: %s' % path) # Split into zipfile and shapefile paths if path.endswith('.zip'): zpath = path @@ -1708,8 +1726,9 @@ def __init__(self, target=None, shapeType=None, autoBalance=False, **kwargs): self.shapeType = shapeType self.shp = self.shx = self.dbf = None if target: + target = pathlike_obj(target) if not is_string(target): - raise Exception('The target filepath {} must be of type str/unicode, not {}.'.format(repr(target), type(target)) ) + raise Exception('The target filepath {} must be of type str/unicode or path-like, not {}.'.format(repr(target), type(target)) ) self.shp = self.__getFileObj(os.path.splitext(target)[0] + '.shp') self.shx = self.__getFileObj(os.path.splitext(target)[0] + '.shx') self.dbf = self.__getFileObj(os.path.splitext(target)[0] + '.dbf') diff --git a/test_shapefile.py b/test_shapefile.py index b08ce38..d175ae2 100644 --- a/test_shapefile.py +++ b/test_shapefile.py @@ -3,15 +3,22 @@ """ # std lib imports import os.path +import sys +if sys.version_info.major == 3: + from pathlib import Path # third party imports import pytest import json import datetime +if sys.version_info.major == 2: + # required by pytest for python <36 + from pathlib2 import Path # our imports import shapefile + # define various test shape tuples of (type, points, parts indexes, and expected geo interface output) geo_interface_tests = [ (shapefile.POINT, # point [(1,1)], @@ -403,6 +410,15 @@ def test_reader_shapefile_extension_ignored(): assert not os.path.exists(filename) +def test_reader_pathlike(): + """ + Assert that path-like objects can be read. + """ + base = Path("shapefiles") + with shapefile.Reader(base / "blockgroups") as sf: + assert len(sf) == 663 + + def test_reader_filelike_dbf_only(): """ Assert that specifying just the @@ -888,6 +904,20 @@ def test_write_default_shp_shx_dbf(tmpdir): assert os.path.exists(filename + ".dbf") +def test_write_pathlike(tmpdir): + """ + Assert that path-like objects can be written. + Similar to test_write_default_shp_shx_dbf. + """ + filename = tmpdir.join("test") + assert not isinstance(filename, str) + with shapefile.Writer(filename) as writer: + writer.field('field1', 'C') + assert (filename + ".shp").ensure() + assert (filename + ".shx").ensure() + assert (filename + ".dbf").ensure() + + def test_write_shapefile_extension_ignored(tmpdir): """ Assert that the filename's extension is