Add auto_mkdir argument to SMBFileSystem (#1604)

Co-authored-by: Martin Durant <[email protected]>

Author: uermel

fsspec/implementations/smb.py

@@ -68,6 +68,7 @@ def __init__(
         encrypt=None,
         share_access=None,
         register_session_retries=5,
+        auto_mkdir=False,
         **kwargs,
     ):
         """
@@ -102,6 +103,10 @@ def __init__(
             - 'r': Allow other handles to be opened with read access.
             - 'w': Allow other handles to be opened with write access.
             - 'd': Allow other handles to be opened with delete access.
+        auto_mkdir: bool
+            Whether, when opening a file, the directory containing it should
+            be created (if it doesn't already exist). This is assumed by pyarrow
+            and zarr-python code.
         """
         super().__init__(**kwargs)
         self.host = host
@@ -113,6 +118,7 @@ def __init__(
         self.temppath = kwargs.pop("temppath", "")
         self.share_access = share_access
         self.register_session_retries = register_session_retries
+        self.auto_mkdir = auto_mkdir
         self._connect()
 
     @property
@@ -224,6 +230,8 @@ def _open(
         By specifying 'share_access' in 'kwargs' it is possible to override the
         default shared access setting applied in the constructor of this object.
         """
+        if self.auto_mkdir and "w" in mode:
+            self.makedirs(self._parent(path), exist_ok=True)
         bls = block_size if block_size is not None and block_size >= 0 else -1
         wpath = _as_unc_path(self.host, path)
         share_access = kwargs.pop("share_access", self.share_access)
@@ -245,6 +253,8 @@ def copy(self, path1, path2, **kwargs):
         """Copy within two locations in the same filesystem"""
         wpath1 = _as_unc_path(self.host, path1)
         wpath2 = _as_unc_path(self.host, path2)
+        if self.auto_mkdir:
+            self.makedirs(self._parent(path2), exist_ok=True)
         smbclient.copyfile(wpath1, wpath2, port=self._port, **kwargs)
 
     def _rm(self, path):

fsspec/implementations/tests/test_smb.py

@@ -89,6 +89,33 @@ def test_simple(smb_params):
     assert not fsmb.exists(adir)
 
 
+@pytest.mark.flaky(reruns=2, reruns_delay=2)
+def test_auto_mkdir(smb_params):
+    adir = "/home/adir"
+    adir2 = "/home/adir/otherdir/"
+    afile = "/home/adir/otherdir/afile"
+    fsmb = fsspec.get_filesystem_class("smb")(**smb_params, auto_mkdir=True)
+    fsmb.touch(afile)
+    assert fsmb.exists(adir)
+    assert fsmb.exists(adir2)
+    assert fsmb.exists(afile)
+    assert fsmb.info(afile)["type"] == "file"
+
+    another_dir = "/home/another_dir"
+    another_dir2 = "/home/another_dir/another_nested_dir/"
+    another_file = "/home/another_dir/another_nested_dir/another_file"
+    fsmb.copy(afile, another_file)
+    assert fsmb.exists(another_dir)
+    assert fsmb.exists(another_dir2)
+    assert fsmb.exists(another_file)
+    assert fsmb.info(another_file)["type"] == "file"
+
+    fsmb.rm(adir, recursive=True)
+    fsmb.rm(another_dir, recursive=True)
+    assert not fsmb.exists(adir)
+    assert not fsmb.exists(another_dir)
+
+
 @pytest.mark.flaky(reruns=2, reruns_delay=2)
 def test_with_url(smb_params):
     if smb_params["port"] is None: