1  import warnings
  2  
  3  from mlflow.utils.docstring_utils import (
  4      ParamDocs,
  5      _indent,
  6      docstring_version_compatibility_warning,
  7      format_docstring,
  8  )
  9  
 10  
 11  def test_indent_empty():
 12      a = ""
 13      b = " " * 4
 14      assert _indent(a, b) == a
 15  
 16  
 17  def test_indent_single_line():
 18      a = "x"
 19      b = " " * 4
 20      assert _indent(a, b) == a
 21  
 22  
 23  def test_indent_multi_line():
 24      a = """x\nx
 25      x\nx
 26      x"""
 27      b = " " * 4
 28      assert _indent(a, b) == "x\n    x\n        x\n    x\n        x"
 29  
 30  
 31  def test_param_docs_format():
 32      pd = ParamDocs({"x": "{{ x }}", "y": "{{ y }}", "z": "{{ x }}, {{ y }}"})
 33      formatted = pd.format(x="a", y="b")
 34      assert isinstance(formatted, ParamDocs)
 35      assert formatted == {"x": "a", "y": "b", "z": "a, b"}
 36  
 37  
 38  def test_param_docs_format_no_changes():
 39      @format_docstring({
 40          "multi_line": """Single line
 41  Another line\n    Another indented line""",
 42          "single_line": "hi",
 43      })
 44      def f(p1, p2, p3, p4):
 45          """asdf
 46  
 47          Args:
 48              p1:
 49                  asdf
 50              p2: asdf
 51              p3:
 52                  asdf
 53              p4:
 54                  asdf
 55          """
 56  
 57      expected = """asdf
 58  
 59          Args:
 60              p1:
 61                  asdf
 62              p2: asdf
 63              p3:
 64                  asdf
 65              p4:
 66                  asdf
 67          """
 68  
 69      assert f.__doc__ == expected
 70      assert f.__name__ == "f"
 71  
 72  
 73  def test_param_docs_format_google():
 74      @format_docstring({
 75          "multi_line": """Single line
 76  Another line\n    Another indented line""",
 77          "single_line": "hi",
 78      })
 79      # fmt: off
 80      def f(p1, p2, p3, p4):
 81          """asdf
 82  
 83          Args:
 84              p1:
 85                  asdf
 86              p2: {{ multi_line }}
 87              p3:
 88                  {{ single_line }}
 89              p4:
 90                  {{ multi_line }}
 91          """
 92  
 93      expected = """asdf
 94  
 95          Args:
 96              p1:
 97                  asdf
 98              p2: Single line
 99                  Another line
100                      Another indented line
101              p3:
102                  hi
103              p4:
104                  Single line
105                  Another line
106                      Another indented line
107          """
108      # fmt: on
109  
110      assert f.__doc__ == expected
111      assert f.__name__ == "f"
112  
113  
114  def test_param_docs_format_not_google():
115      @format_docstring({
116          "multi_line": """Single line
117  Another line\n    Another indented line""",
118          "single_line": "hi",
119      })
120      # fmt: off
121      def f():
122          """
123          asdf
124  
125          Args
126              p1: asdf
127              p2: {{ multi_line }}
128              p3: {{ single_line }}
129              p4:
130                  {{ multi_line }}
131          """
132  
133      expected = """
134          asdf
135  
136          Args
137              p1: asdf
138              p2: Single line
139                  Another line
140                      Another indented line
141              p3: hi
142              p4:
143                  Single line
144                  Another line
145                      Another indented line
146          """
147      # fmt: on
148  
149      assert f.__doc__ == expected
150      assert f.__name__ == "f"
151  
152  
153  def test_docstring_version_compatibility_warning():
154      @docstring_version_compatibility_warning("sklearn")
155      def func():
156          pass
157  
158      @docstring_version_compatibility_warning("sklearn")
159      def another_func():
160          func()
161  
162      with warnings.catch_warnings(record=True) as w:
163          func()
164  
165      # Exclude irrelevant warnings
166      warns = [x for x in w if "MLflow Models integration is known to be compatible" in str(x)]
167      assert len(warns) == 0
168  
169      with warnings.catch_warnings(record=True) as w:
170          another_func()
171  
172      warns = [x for x in w if "MLflow Models integration is known to be compatible" in str(x)]
173      assert len(warns) == 0