Repository: IDSIA/sacred (python)

Issue:
Sacred dislikes return type hint in config
I have a config function with a return type hint, like that:
```python
@ex.config
def config() -> None:
    """Experiment config."""
    ...
```
It seems that Sacred is trying to match this against an overly restrictive regex, and it gives an assert error:
```
Traceback (most recent call last):
  File "/home/carlos/git/pytorch_congealing/main.py", line 11, in <module>
    def config() -> None:
  File "/home/carlos/Programas/Utilidades/Lenguajes/miniconda3/envs/fda/lib/python3.7/site-packages/sacred/ingredient.py", line 162, in config
    self.configurations.append(ConfigScope(function))
  File "/home/carlos/Programas/Utilidades/Lenguajes/miniconda3/envs/fda/lib/python3.7/site-packages/sacred/config/config_scope.py", line 26, in __init__
    self._body_code = get_function_body_code(func)
  File "/home/carlos/Programas/Utilidades/Lenguajes/miniconda3/envs/fda/lib/python3.7/site-packages/sacred/config/config_scope.py", line 149, in get_function_body_code
    func_body, line_offset = get_function_body(func)
  File "/home/carlos/Programas/Utilidades/Lenguajes/miniconda3/envs/fda/lib/python3.7/site-packages/sacred/config/config_scope.py", line 107, in get_function_body
    assert defs
AssertionError
```

Pull request description:
Allow config scopes with type annotations.
Allows the usage of type hints in config scopes. Closes #818.

Notes:
- A new test is added for this behaviour.
- Instead of fixing the regexp, it uses `ast` and `tokenize` modules. With `ast` we find the first line with code in the function body. As comments are removed in the AST representation, we then generate the tokens of the function until that line is reached, in order to keep the previous lines that only have comments and whitespace.

Patch:
diff --git a/.gitignore b/.gitignore
index 213dd8a..815a545 100644
--- a/.gitignore
+++ b/.gitignore
@@ -93,4 +93,5 @@ com_crashlytics_export_strings.xml
 # GEdit temporary files
 *~
 
+/.pytest_cache/
 pip-wheel-metadata/
diff --git a/sacred/config/config_scope.py b/sacred/config/config_scope.py
index 631aae2..0ca478f 100644
--- a/sacred/config/config_scope.py
+++ b/sacred/config/config_scope.py
@@ -5,14 +5,16 @@ import ast
 import inspect
 import io
 import re
-from tokenize import tokenize, TokenError, COMMENT
-from copy import copy
+import textwrap
+import token
 
+from copy import copy
 from sacred import SETTINGS
 from sacred.config.config_summary import ConfigSummary
 from sacred.config.utils import dogmatize, normalize_or_die, recursive_fill_in
 from sacred.config.signature import get_argspec
 from sacred.utils import ConfigError
+from tokenize import generate_tokens, tokenize, TokenError, COMMENT
 
 
 class ConfigScope:
@@ -94,20 +96,53 @@ class ConfigScope:
 
 def get_function_body(func):
     func_code_lines, start_idx = inspect.getsourcelines(func)
-    func_code = "".join(func_code_lines)
-    arg = "(?:[a-zA-Z_][a-zA-Z0-9_]*)"
-    arguments = r"{0}(?:\s*,\s*{0})*,?".format(arg)
-    func_def = re.compile(
-        r"^[ \t]*def[ \t]*{}[ \t]*\(\s*({})?\s*\)[ \t]*:[ \t]*(?:#[^\n]*)?\n".format(
-            func.__name__, arguments
-        ),
-        flags=re.MULTILINE,
+    func_code = textwrap.dedent("".join(func_code_lines))
+    # Lines are now dedented
+    func_code_lines = func_code.splitlines(True)
+    func_ast = ast.parse(func_code)
+    first_code = func_ast.body[0].body[0]
+    line_offset = first_code.lineno
+    col_offset = first_code.col_offset
+
+    # Add also previous empty / comment lines
+    acceptable_tokens = {
+        token.NEWLINE,
+        token.INDENT,
+        token.DEDENT,
+        token.COMMENT,
+        token.ENDMARKER,
+    }
+    last_token_type_acceptable = True
+    line_offset_fixed = line_offset
+    col_offset_fixed = col_offset
+    iterator = iter(func_code_lines)
+    for parsed_token in generate_tokens(lambda: next(iterator)):
+
+        token_acceptable = parsed_token.type in acceptable_tokens or (
+            parsed_token.type == token.NL and last_token_type_acceptable
+        )
+
+        # If the token ends after the start of the first code,
+        # we have finished
+        if parsed_token.end[0] > line_offset or (
+            parsed_token.end[0] == line_offset and parsed_token.end[1] >= col_offset
+        ):
+            break
+
+        if not token_acceptable:
+            line_offset_fixed = parsed_token.end[0]
+            col_offset_fixed = parsed_token.end[1]
+
+        last_token_type_acceptable = token_acceptable
+
+    func_body = (
+        # First line, without first part if needed
+        func_code_lines[line_offset_fixed - 1][col_offset_fixed:]
+        # Rest of the lines
+        + "".join(func_code_lines[line_offset_fixed:])
     )
-    defs = list(re.finditer(func_def, func_code))
-    assert defs
-    line_offset = start_idx + func_code[: defs[0].end()].count("\n") - 1
-    func_body = func_code[defs[0].end() :]
-    return func_body, line_offset
+
+    return func_body, start_idx + line_offset_fixed
 
 
 def is_empty_or_comment(line):


Test patch:
diff --git a/tests/test_config/test_config_scope.py b/tests/test_config/test_config_scope.py
index 30c6e44..44e6e06 100644
--- a/tests/test_config/test_config_scope.py
+++ b/tests/test_config/test_config_scope.py
@@ -318,7 +318,8 @@ strings
         def subfunc():
             return 23
 
-body = '''# Lets do the most evil things with indentation
+body = '''  # test comment
+# Lets do the most evil things with indentation
   # 1
     # 2
      # ran
@@ -343,7 +344,8 @@ strings
             return 23
 '''
 
-dedented_body = '''# Lets do the most evil things with indentation
+dedented_body = '''# test comment
+# Lets do the most evil things with indentation
 # 1
 # 2
 # ran
@@ -382,3 +384,22 @@ def test_get_function_body():
 def test_config_scope_can_deal_with_indentation_madness():
     # assert_no_raise:
     ConfigScope(evil_indentation_func)
+
+
+def test_conf_scope_with_type_info():
+    @ConfigScope
+    def conf_scope(a: int) -> None:
+        answer = 2 * a
+
+    cfg = conf_scope(preset={"a": 21})
+    assert cfg["answer"] == 42
+
+
+def test_conf_scope_in_same_line():
+    # fmt: off
+    @ConfigScope
+    def conf_scope(a: int) -> None: answer = 2 * a
+    # fmt: on
+
+    cfg = conf_scope(preset={"a": 21})
+    assert cfg["answer"] == 42