Only handle TOKEN.NUMBER and TOKEN.STRING

Signed-off-by: Pedro Algarvio <[email protected]>
saltstack · Jun 18, 2023 · 77eb5c5 · 77eb5c5
1 parent a75e0d0
commit 77eb5c5
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 0 deletions.
diff --git a/src/saltrewrite/salt/fix_warn_until.py b/src/saltrewrite/salt/fix_warn_until.py
@@ -37,11 +37,28 @@ def rewrite(paths, interactive=False, silent=False):
             )
             """
         )
+        .filter(filter_versions)
         .modify(fix_warn_unil_version)
         .execute(write=True, interactive=interactive, silent=silent)
     )
 
 
+def filter_versions(node, capture, filename):
+    """
+    If the first child is a docstring, process it
+    """
+    arglist = capture["function_arguments"][0]
+    warn_until_version_argument = arglist.children[0]
+    if warn_until_version_argument.type == TOKEN.NUMBER:
+        # For example, 3008, 3009.0
+        return True
+    if warn_until_version_argument.type == TOKEN.STRING:
+        # For example, "3008", "Aragon"
+        return True
+    # Don't process node
+    return False
+
+
 def fix_warn_unil_version(node, capture, filename):
     """
     Automaticaly run fixes against docstrings

diff --git a/tests/salt/test_warn_until.py b/tests/salt/test_warn_until.py
@@ -20,6 +20,10 @@ def two():
         def three():
             warn_until(3008.0, "Code deprecated in 3008.0")
             print("one")
+
+        def four():
+            warn_until("3008", "Code deprecated in 3008.0")
+            print("one")
         """
     )
     expected_code = textwrap.dedent(
@@ -37,6 +41,10 @@ def two():
         def three():
             warn_until(3008.0, "Code deprecated in 3008.0")
             print("one")
+
+        def four():
+            warn_until(3008.0, "Code deprecated in 3008.0")
+            print("one")
         """
     )
     fpath = tempfiles.makepyfile(code, prefix="test_")
@@ -62,6 +70,10 @@ def two():
         def three():
             salt.utils.versions.warn_until(3008.0, "Code deprecated in 3008.0")
             print("one")
+
+        def four():
+            salt.utils.versions.warn_until("3008", "Code deprecated in 3008.0")
+            print("one")
         """
     )
     expected_code = textwrap.dedent(
@@ -79,6 +91,10 @@ def two():
         def three():
             salt.utils.versions.warn_until(3008.0, "Code deprecated in 3008.0")
             print("one")
+
+        def four():
+            salt.utils.versions.warn_until(3008.0, "Code deprecated in 3008.0")
+            print("one")
         """
     )
     fpath = tempfiles.makepyfile(code, prefix="test_")