Fix tilde expansion for LocalConnector.put_file()

pyinfra/connectors/local.py

@@ -1,4 +1,5 @@
 import os
+import shlex
 from shutil import which
 from tempfile import mkstemp
 from typing import TYPE_CHECKING, Tuple
@@ -118,6 +119,12 @@ def put_file(
             bool: Indicating success or failure
         """
 
+        if remote_filename.startswith("~/"):
+            # Do not quote leading tilde to ensure that it gets properly expanded by the shell
+            remote_filename = f"~/{shlex.quote(remote_filename[2:])}"
+        else:
+            remote_filename = QuoteString(remote_filename)
+
         _, temp_filename = mkstemp()
 
         try:
@@ -133,7 +140,7 @@ def put_file(
 
             # Copy the file using `cp` such that we support sudo/su
             status, output = self.run_shell_command(
-                StringCommand("cp", temp_filename, QuoteString(remote_filename)),
+                StringCommand("cp", temp_filename, remote_filename),
                 print_output=print_output,
                 print_input=print_input,
                 **arguments,
@@ -146,6 +153,7 @@ def put_file(
 
         if print_output:
             click.echo(
+                # TODO: Check if the modification of remote_filename affects the output
                 "{0}file copied: {1}".format(self.host.print_prefix, remote_filename),
                 err=True,
             )

tests/test_connectors/test_local.py

@@ -158,6 +158,44 @@ def test_put_file_with_spaces(self):
             stdin=PIPE,
         )
 
+    def test_put_file_tilde_expansion(self):
+        inventory = make_inventory(hosts=("@local",))
+        State(inventory, Config())
+
+        host = inventory.get_host("@local")
+
+        fake_process = MagicMock(returncode=0)
+        self.fake_popen_mock.return_value = fake_process
+
+        host.put_file("not-a-file", "~/file-in-home-directory", print_output=True)
+
+        self.fake_popen_mock.assert_called_with(
+            "sh -c 'cp __tempfile__ ~/file-in-home-directory'",
+            shell=True,
+            stdout=PIPE,
+            stderr=PIPE,
+            stdin=PIPE,
+        )
+
+    def test_put_file_tilde_expansion_with_spaces(self):
+        inventory = make_inventory(hosts=("@local",))
+        State(inventory, Config())
+
+        host = inventory.get_host("@local")
+
+        fake_process = MagicMock(returncode=0)
+        self.fake_popen_mock.return_value = fake_process
+
+        host.put_file("not-a-file", "~/not another file with spaces", print_output=True)
+
+        self.fake_popen_mock.assert_called_with(
+            "sh -c 'cp __tempfile__ ~/'\"'\"'not another file with spaces'\"'\"''",
+            shell=True,
+            stdout=PIPE,
+            stderr=PIPE,
+            stdin=PIPE,
+        )
+
     def test_put_file_error(self):
         inventory = make_inventory(hosts=("@local",))
         State(inventory, Config())