From 3346a013070859bf6ea64e04acdcfc8eee7f09f0 Mon Sep 17 00:00:00 2001
From: skshetry <18718008+skshetry@users.noreply.github.com>
Date: Wed, 6 Dec 2023 08:55:02 +0545
Subject: [PATCH] microopt: reorder to remove assertions, lazy imports (#237)

---
 src/dvc_objects/fs/system.py | 55 ++++++++++++++----------------------
 tests/benchmarks/test_fs.py  |  9 ++++--
 tests/test_reflink.py        |  3 +-
 3 files changed, 29 insertions(+), 38 deletions(-)

diff --git a/src/dvc_objects/fs/system.py b/src/dvc_objects/fs/system.py
index 1e06cf7..5becbd2 100644
--- a/src/dvc_objects/fs/system.py
+++ b/src/dvc_objects/fs/system.py
@@ -2,10 +2,8 @@
 import errno
 import logging
 import os
-import platform
 import stat
 import sys
-import functools
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -34,7 +32,6 @@ def symlink(source: "AnyFSPath", link_name: "AnyFSPath") -> None:
     os.symlink(source, link_name)
 
 
-@functools.lru_cache(maxsize=1)
 def _clonefile():
     def _cdll(name):
         return ctypes.CDLL(name, use_errno=True)
@@ -68,53 +65,43 @@ def _cdll(name):
     return clonefile
 
 
-def _reflink_darwin(src: "AnyFSPath", dst: "AnyFSPath") -> None:
+# NOTE: reflink may (macos) or may not (linux) clone permissions,
+# so the user needs to handle those himself.
+if sys.platform == "darwin":
     clonefile = _clonefile()
 
-    ret = clonefile(
-        ctypes.c_char_p(os.fsencode(src)),
-        ctypes.c_char_p(os.fsencode(dst)),
-        ctypes.c_int(0),
-    )
-    if ret:
-        err = ctypes.get_errno()
-        msg = os.strerror(err)
-        raise OSError(err, msg)
-
+    def reflink(src, dst):
+        ret = clonefile(
+            ctypes.c_char_p(os.fsencode(src)),
+            ctypes.c_char_p(os.fsencode(dst)),
+            ctypes.c_int(0),
+        )
+        if ret:
+            err = ctypes.get_errno()
+            msg = os.strerror(err)
+            raise OSError(err, msg)
 
-def _reflink_linux(src: "AnyFSPath", dst: "AnyFSPath") -> None:
+elif sys.platform == "linux":
     import fcntl  # pylint: disable=import-error
-    from contextlib import suppress
 
     FICLONE = 0x40049409
 
-    if sys.platform != "linux":
-        raise AssertionError
-    else:
+    def reflink(src, dst):
         try:
             with open(src, "rb") as s, open(dst, "wb+") as d:
                 fcntl.ioctl(d.fileno(), FICLONE, s.fileno())
         except OSError:
-            with suppress(OSError):
+            try:
                 os.unlink(dst)
+            except OSError:
+                pass
             raise
+else:
 
-
-def reflink(source: "AnyFSPath", link_name: "AnyFSPath") -> None:
-    source, link_name = os.fspath(source), os.fspath(link_name)
-
-    # NOTE: reflink may (macos) or may not (linux) clone permissions,
-    # so the user needs to handle those himself.
-
-    system = platform.system()
-    if system == "Darwin":
-        _reflink_darwin(source, link_name)
-    elif system == "Linux":
-        _reflink_linux(source, link_name)
-    else:
+    def reflink(src, dst):
         raise OSError(
             errno.ENOTSUP,
-            f"reflink is not supported on '{system}'",
+            f"reflink is not supported on '{sys.platform}'",
         )
 
 
diff --git a/tests/benchmarks/test_fs.py b/tests/benchmarks/test_fs.py
index 58e5aea..02df746 100644
--- a/tests/benchmarks/test_fs.py
+++ b/tests/benchmarks/test_fs.py
@@ -1,15 +1,18 @@
 import errno
-import pytest
 import shutil
+
+import pytest
 from reflink import reflink as pyreflink
 from reflink.error import ReflinkImpossibleError
-from dvc_objects.fs.system import reflink, hardlink, symlink
+
+from dvc_objects.fs.system import hardlink, reflink, symlink
 
 NLINKS = 10000
 
 
 @pytest.mark.parametrize(
-    "link", [pytest.param(pyreflink, id="pyreflink"), reflink, hardlink, symlink]
+    "link",
+    [pytest.param(pyreflink, id="pyreflink"), reflink, hardlink, symlink],
 )
 def test_link(benchmark, tmp_path, link):
     (tmp_path / "original").mkdir()
diff --git a/tests/test_reflink.py b/tests/test_reflink.py
index b8d736f..2768315 100644
--- a/tests/test_reflink.py
+++ b/tests/test_reflink.py
@@ -1,4 +1,5 @@
 import errno
+import os
 import stat
 from os import fspath, umask
 from pathlib import Path
@@ -37,12 +38,12 @@ def test_reflink(test_dir):
     assert stat_mode == (0o666 & ~umask(0))
 
 
+@pytest.mark.skipif(os.name != "nt", reason="only run in Windows")
 def test_reflink_unsupported_on_windows(test_dir, mocker):
     src = test_dir / "source"
     dest = test_dir / "dest"
     src.write_bytes(b"content")
 
-    mocker.patch("platform.system", mocker.MagicMock(return_value="Windows"))
     with pytest.raises(OSError) as exc:
         reflink(fspath(src), fspath(dest))