tests_fetcher.py 46.5 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# coding=utf-8
# Copyright 2021 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Sylvain Gugger's avatar
Sylvain Gugger committed
16
17
"""
Welcome to tests_fetcher V2.
Sylvain Gugger's avatar
Sylvain Gugger committed
18

Sylvain Gugger's avatar
Sylvain Gugger committed
19
20
21
This util is designed to fetch tests to run on a PR so that only the tests impacted by the modifications are run, and
when too many models are being impacted, only run the tests of a subset of core models. It works like this.

22
23
24
Stage 1: Identify the modified files. For jobs that run on the main branch, it's just the diff with the last commit.
On a PR, this takes all the files from the branching point to the current commit (so all modifications in a PR, not
just the last commit) but excludes modifications that are on docstrings or comments only.
Sylvain Gugger's avatar
Sylvain Gugger committed
25
26
27
28

Stage 2: Extract the tests to run. This is done by looking at the imports in each module and test file: if module A
imports module B, then changing module B impacts module A, so the tests using module A should be run. We thus get the
dependencies of each model and then recursively builds the 'reverse' map of dependencies to get all modules and tests
29
impacted by a given file. We then only keep the tests (and only the core models tests if there are too many modules).
Sylvain Gugger's avatar
Sylvain Gugger committed
30
31
32
33
34
35

Caveats:
  - This module only filters tests by files (not individual tests) so it's better to have tests for different things
    in different files.
  - This module assumes inits are just importing things, not really building objects, so it's better to structure
    them this way and move objects building in separate submodules.
36
37
38
39
40
41
42
43
44
45
46
47
48
49

Usage:

Base use to fetch the tests in a pull request

```bash
python utils/tests_fetcher.py
```

Base use to fetch the tests on a the main branch (with diff from the last commit):

```bash
python utils/tests_fetcher.py --diff_with_last_commit
```
Sylvain Gugger's avatar
Sylvain Gugger committed
50
51
"""

52
53
import argparse
import collections
Yih-Dar's avatar
Yih-Dar committed
54
import json
55
56
57
58
import os
import re
from contextlib import contextmanager
from pathlib import Path
59
from typing import Dict, List, Optional, Tuple, Union
60
61
62
63

from git import Repo


Sylvain Gugger's avatar
Sylvain Gugger committed
64
PATH_TO_REPO = Path(__file__).parent.parent.resolve()
65
PATH_TO_EXAMPLES = PATH_TO_REPO / "examples"
Sylvain Gugger's avatar
Sylvain Gugger committed
66
67
68
69
70
PATH_TO_TRANFORMERS = PATH_TO_REPO / "src/transformers"
PATH_TO_TESTS = PATH_TO_REPO / "tests"

# List here the models to always test.
IMPORTANT_MODELS = [
71
    "auto",
Sylvain Gugger's avatar
Sylvain Gugger committed
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
    # Most downloaded models
    "bert",
    "clip",
    "t5",
    "xlm-roberta",
    "gpt2",
    "bart",
    "mpnet",
    "gpt-j",
    "wav2vec2",
    "deberta-v2",
    "layoutlm",
    "opt",
    "longformer",
    "vit",
    # Pipeline-specific model (to be sure each pipeline has one model in this list)
    "tapas",
    "vilt",
    "clap",
    "detr",
    "owlvit",
    "dpt",
    "videomae",
]
96

97
98

@contextmanager
99
def checkout_commit(repo: Repo, commit_id: str):
100
    """
101
102
103
104
105
    Context manager that checks out a given commit when entered, but gets back to the reference it was at on exit.

    Args:
        repo (`git.Repo`): A git repository (for instance the Transformers repo).
        commit_id (`str`): The commit reference to checkout inside the context manager.
106
107
108
109
110
111
112
113
114
115
116
    """
    current_head = repo.head.commit if repo.head.is_detached else repo.head.ref

    try:
        repo.git.checkout(commit_id)
        yield

    finally:
        repo.git.checkout(current_head)


117
def clean_code(content: str) -> str:
118
    """
119
120
121
122
123
124
125
126
    Remove docstrings, empty line or comments from some code (used to detect if a diff is real or only concern
    comments or docstings).

    Args:
        content (`str`): The code to clean

    Returns:
        `str`: The cleaned code.
127
    """
128
129
    # We need to deactivate autoformatting here to write escaped triple quotes (we cannot use real triple quotes or
    # this would mess up the result if this function applied to this particular file).
130
131
132
133
134
135
136
137
138
139
140
141
142
    # fmt: off
    # Remove docstrings by splitting on triple " then triple ':
    splits = content.split('\"\"\"')
    content = "".join(splits[::2])
    splits = content.split("\'\'\'")
    # fmt: on
    content = "".join(splits[::2])

    # Remove empty lines and comments
    lines_to_keep = []
    for line in content.split("\n"):
        # remove anything that is after a # sign.
        line = re.sub("#.*$", "", line)
143
144
145
        # remove white lines
        if len(line) != 0 and not line.isspace():
            lines_to_keep.append(line)
146
147
148
    return "\n".join(lines_to_keep)


149
def keep_doc_examples_only(content: str) -> str:
150
    """
151
152
153
154
155
156
157
158
    Remove everything from the code content except the doc examples (used to determined if a diff should trigger doc
    tests or not).

    Args:
        content (`str`): The code to clean

    Returns:
        `str`: The cleaned code.
159
160
161
162
163
164
165
166
167
168
169
    """
    # Keep doc examples only by splitting on triple "`"
    splits = content.split("```")
    # Add leading and trailing "```" so the navigation is easier when compared to the original input `content`
    content = "```" + "```".join(splits[1::2]) + "```"

    # Remove empty lines and comments
    lines_to_keep = []
    for line in content.split("\n"):
        # remove anything that is after a # sign.
        line = re.sub("#.*$", "", line)
170
171
172
        # remove white lines
        if len(line) != 0 and not line.isspace():
            lines_to_keep.append(line)
173
174
175
    return "\n".join(lines_to_keep)


176
def get_all_tests() -> List[str]:
Yih-Dar's avatar
Yih-Dar committed
177
    """
178
179
    Walks the `tests` folder to return a list of files/subfolders. This is used to split the tests to run when using
    paralellism. The split is:
Yih-Dar's avatar
Yih-Dar committed
180

181
    - folders under `tests`: (`tokenization`, `pipelines`, etc) except the subfolder `models` is excluded.
Yih-Dar's avatar
Yih-Dar committed
182
183
184
185
186
    - folders under `tests/models`: `bert`, `gpt2`, etc.
    - test files under `tests`: `test_modeling_common.py`, `test_tokenization_common.py`, etc.
    """

    # test folders/files directly under `tests` folder
Sylvain Gugger's avatar
Sylvain Gugger committed
187
188
189
    tests = os.listdir(PATH_TO_TESTS)
    tests = [f"tests/{f}" for f in tests if "__pycache__" not in f]
    tests = sorted([f for f in tests if (PATH_TO_REPO / f).is_dir() or f.startswith("tests/test_")])
Yih-Dar's avatar
Yih-Dar committed
190
191

    # model specific test folders
Sylvain Gugger's avatar
Sylvain Gugger committed
192
193
194
    model_test_folders = os.listdir(PATH_TO_TESTS / "models")
    model_test_folders = [f"tests/models/{f}" for f in model_test_folders if "__pycache__" not in f]
    model_test_folders = sorted([f for f in model_test_folders if (PATH_TO_REPO / f).is_dir()])
Yih-Dar's avatar
Yih-Dar committed
195
196

    tests.remove("tests/models")
Sylvain Gugger's avatar
Sylvain Gugger committed
197
198
199
    # Sagemaker tests are not meant to be run on the CI.
    if "tests/sagemaker" in tests:
        tests.remove("tests/sagemaker")
Yih-Dar's avatar
Yih-Dar committed
200
201
202
203
204
    tests = model_test_folders + tests

    return tests


205
def diff_is_docstring_only(repo: Repo, branching_point: str, filename: str) -> bool:
206
    """
207
208
209
210
211
212
213
214
215
    Check if the diff is only in docstrings (or comments and whitespace) in a filename.

    Args:
        repo (`git.Repo`): A git repository (for instance the Transformers repo).
        branching_point (`str`): The commit reference of where to compare for the diff.
        filename (`str`): The filename where we want to know if the diff isonly in docstrings/comments.

    Returns:
        `bool`: Whether the diff is docstring/comments only or not.
216
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
217
    folder = Path(repo.working_dir)
218
    with checkout_commit(repo, branching_point):
Sylvain Gugger's avatar
Sylvain Gugger committed
219
        with open(folder / filename, "r", encoding="utf-8") as f:
220
221
            old_content = f.read()

Sylvain Gugger's avatar
Sylvain Gugger committed
222
    with open(folder / filename, "r", encoding="utf-8") as f:
223
224
225
226
227
228
229
230
        new_content = f.read()

    old_content_clean = clean_code(old_content)
    new_content_clean = clean_code(new_content)

    return old_content_clean == new_content_clean


231
def diff_contains_doc_examples(repo: Repo, branching_point: str, filename: str) -> bool:
232
    """
233
234
235
236
237
238
239
240
241
    Check if the diff is only in code examples of the doc in a filename.

    Args:
        repo (`git.Repo`): A git repository (for instance the Transformers repo).
        branching_point (`str`): The commit reference of where to compare for the diff.
        filename (`str`): The filename where we want to know if the diff is only in codes examples.

    Returns:
        `bool`: Whether the diff is only in code examples of the doc or not.
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
    """
    folder = Path(repo.working_dir)
    with checkout_commit(repo, branching_point):
        with open(folder / filename, "r", encoding="utf-8") as f:
            old_content = f.read()

    with open(folder / filename, "r", encoding="utf-8") as f:
        new_content = f.read()

    old_content_clean = keep_doc_examples_only(old_content)
    new_content_clean = keep_doc_examples_only(new_content)

    return old_content_clean != new_content_clean


257
def get_diff(repo: Repo, base_commit: str, commits: List[str]) -> List[str]:
258
    """
259
260
261
262
263
264
265
266
267
268
269
270
271
272
    Get the diff between a base commit and one or several commits.

    Args:
        repo (`git.Repo`):
            A git repository (for instance the Transformers repo).
        base_commit (`str`):
            The commit reference of where to compare for the diff. This is the current commit, not the branching point!
        commits (`List[str]`):
            The list of commits with which to compare the repo at `base_commit` (so the branching point).

    Returns:
        `List[str]`: The list of Python files with a diff (files added, renamed or deleted are always returned, files
        modified are returned if the diff in the file is not only in docstrings or comments, see
        `diff_is_docstring_only`).
273
    """
274
275
    print("\n### DIFF ###\n")
    code_diff = []
276
277
    for commit in commits:
        for diff_obj in commit.diff(base_commit):
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
            # We always add new python files
            if diff_obj.change_type == "A" and diff_obj.b_path.endswith(".py"):
                code_diff.append(diff_obj.b_path)
            # We check that deleted python files won't break corresponding tests.
            elif diff_obj.change_type == "D" and diff_obj.a_path.endswith(".py"):
                code_diff.append(diff_obj.a_path)
            # Now for modified files
            elif diff_obj.change_type in ["M", "R"] and diff_obj.b_path.endswith(".py"):
                # In case of renames, we'll look at the tests using both the old and new name.
                if diff_obj.a_path != diff_obj.b_path:
                    code_diff.extend([diff_obj.a_path, diff_obj.b_path])
                else:
                    # Otherwise, we check modifications are in code and not docstrings.
                    if diff_is_docstring_only(repo, commit, diff_obj.b_path):
                        print(f"Ignoring diff in {diff_obj.b_path} as it only concerns docstrings or comments.")
                    else:
                        code_diff.append(diff_obj.a_path)

    return code_diff


299
def get_modified_python_files(diff_with_last_commit: bool = False) -> List[str]:
Sylvain Gugger's avatar
Sylvain Gugger committed
300
301
302
303
304
    """
    Return a list of python files that have been modified between:

    - the current head and the main branch if `diff_with_last_commit=False` (default)
    - the current head and its parent commit otherwise.
305
306
307
308
309

    Returns:
        `List[str]`: The list of Python files with a diff (files added, renamed or deleted are always returned, files
        modified are returned if the diff in the file is not only in docstrings or comments, see
        `diff_is_docstring_only`).
Sylvain Gugger's avatar
Sylvain Gugger committed
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
    """
    repo = Repo(PATH_TO_REPO)

    if not diff_with_last_commit:
        print(f"main is at {repo.refs.main.commit}")
        print(f"Current head is at {repo.head.commit}")

        branching_commits = repo.merge_base(repo.refs.main, repo.head)
        for commit in branching_commits:
            print(f"Branching commit: {commit}")
        return get_diff(repo, repo.head.commit, branching_commits)
    else:
        print(f"main is at {repo.head.commit}")
        parent_commits = repo.head.commit.parents
        for commit in parent_commits:
            print(f"Parent commit: {commit}")
        return get_diff(repo, repo.head.commit, parent_commits)


329
def get_diff_for_doctesting(repo: Repo, base_commit: str, commits: List[str]) -> List[str]:
330
    """
331
332
333
334
335
336
337
338
339
340
341
342
343
    Get the diff in doc examples between a base commit and one or several commits.

    Args:
        repo (`git.Repo`):
            A git repository (for instance the Transformers repo).
        base_commit (`str`):
            The commit reference of where to compare for the diff. This is the current commit, not the branching point!
        commits (`List[str]`):
            The list of commits with which to compare the repo at `base_commit` (so the branching point).

    Returns:
        `List[str]`: The list of Python and Markdown files with a diff (files added or renamed are always returned, files
        modified are returned if the diff in the file is only in doctest examples).
344
345
346
347
348
    """
    print("\n### DIFF ###\n")
    code_diff = []
    for commit in commits:
        for diff_obj in commit.diff(base_commit):
349
350
351
            # We only consider Python files and doc files.
            if not diff_obj.b_path.endswith(".py") and not diff_obj.b_path.endswith(".md"):
                continue
352
            # We always add new python/md files
353
            if diff_obj.change_type in ["A"]:
354
                code_diff.append(diff_obj.b_path)
355
            # Now for modified files
356
            elif diff_obj.change_type in ["M", "R"]:
357
358
359
360
361
362
363
364
365
                # In case of renames, we'll look at the tests using both the old and new name.
                if diff_obj.a_path != diff_obj.b_path:
                    code_diff.extend([diff_obj.a_path, diff_obj.b_path])
                else:
                    # Otherwise, we check modifications contain some doc example(s).
                    if diff_contains_doc_examples(repo, commit, diff_obj.b_path):
                        code_diff.append(diff_obj.a_path)
                    else:
                        print(f"Ignoring diff in {diff_obj.b_path} as it doesn't contain any doc example.")
366
367
368
369

    return code_diff


370
def get_doctest_files(diff_with_last_commit: bool = False) -> List[str]:
371
    """
372
    Return a list of python and Markdown files where doc example have been modified between:
373
374
375

    - the current head and the main branch if `diff_with_last_commit=False` (default)
    - the current head and its parent commit otherwise.
376
377
378
379

    Returns:
        `List[str]`: The list of Python and Markdown files with a diff (files added or renamed are always returned, files
        modified are returned if the diff in the file is only in doctest examples).
380
381
382
    """
    repo = Repo(PATH_TO_REPO)

383
    test_files_to_run = []  # noqa
384
385
386
387
388
389
390
    if not diff_with_last_commit:
        print(f"main is at {repo.refs.main.commit}")
        print(f"Current head is at {repo.head.commit}")

        branching_commits = repo.merge_base(repo.refs.main, repo.head)
        for commit in branching_commits:
            print(f"Branching commit: {commit}")
391
        test_files_to_run = get_diff_for_doctesting(repo, repo.head.commit, branching_commits)
392
393
394
395
396
    else:
        print(f"main is at {repo.head.commit}")
        parent_commits = repo.head.commit.parents
        for commit in parent_commits:
            print(f"Parent commit: {commit}")
397
        test_files_to_run = get_diff_for_doctesting(repo, repo.head.commit, parent_commits)
398

Yih-Dar's avatar
Yih-Dar committed
399
400
401
    # These are files not doctested yet.
    with open("utils/not_doctested.txt") as fp:
        not_doctested = set(fp.read().strip().split("\n"))
402
    # Do not run slow doctest tests
403
404
405
    with open("utils/slow_documentation_tests.txt") as fp:
        slow_documentation_tests = set(fp.read().strip().split("\n"))

406
    # So far we don't have 100% coverage for doctest. This line will be removed once we achieve 100%.
Yih-Dar's avatar
Yih-Dar committed
407
408
409
410
411
412
413
414
    test_files_to_run = [x for x in test_files_to_run if x not in not_doctested and x not in slow_documentation_tests]

    # The file `utils/not_doctested.txt` doesn't contain all files that are not doc-tested, so we need more filters.
    # 1. only include files in `src` or `docs/source/en/`
    test_files_to_run = [x for x in test_files_to_run if x.startswith(("src/", "docs/source/en/"))]
    # 2. not include init files
    test_files_to_run = [x for x in test_files_to_run if not x.endswith(("__init__.py",))]

415
416
417
418
419
420
    # Make sure we did not end up with a test file that was removed
    test_files_to_run = [f for f in test_files_to_run if (PATH_TO_REPO / f).exists()]

    return test_files_to_run


Sylvain Gugger's avatar
Sylvain Gugger committed
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
# (:?^|\n) -> Non-catching group for the beginning of the doc or a new line.
# \s*from\s+(\.+\S+)\s+import\s+([^\n]+) -> Line only contains from .xxx import yyy and we catch .xxx and yyy
# (?=\n) -> Look-ahead to a new line. We can't just put \n here or using find_all on this re will only catch every
#           other import.
_re_single_line_relative_imports = re.compile(r"(?:^|\n)\s*from\s+(\.+\S+)\s+import\s+([^\n]+)(?=\n)")
# (:?^|\n) -> Non-catching group for the beginning of the doc or a new line.
# \s*from\s+(\.+\S+)\s+import\s+\(([^\)]+)\) -> Line continues with from .xxx import (yyy) and we catch .xxx and yyy
# yyy will take multiple lines otherwise there wouldn't be parenthesis.
_re_multi_line_relative_imports = re.compile(r"(?:^|\n)\s*from\s+(\.+\S+)\s+import\s+\(([^\)]+)\)")
# (:?^|\n) -> Non-catching group for the beginning of the doc or a new line.
# \s*from\s+transformers(\S*)\s+import\s+([^\n]+) -> Line only contains from transformers.xxx import yyy and we catch
#           .xxx and yyy
# (?=\n) -> Look-ahead to a new line. We can't just put \n here or using find_all on this re will only catch every
#           other import.
_re_single_line_direct_imports = re.compile(r"(?:^|\n)\s*from\s+transformers(\S*)\s+import\s+([^\n]+)(?=\n)")
# (:?^|\n) -> Non-catching group for the beginning of the doc or a new line.
# \s*from\s+transformers(\S*)\s+import\s+\(([^\)]+)\) -> Line continues with from transformers.xxx import (yyy) and we
# catch .xxx and yyy. yyy will take multiple lines otherwise there wouldn't be parenthesis.
_re_multi_line_direct_imports = re.compile(r"(?:^|\n)\s*from\s+transformers(\S*)\s+import\s+\(([^\)]+)\)")


442
def extract_imports(module_fname: str, cache: Dict[str, List[str]] = None) -> List[str]:
443
    """
444
445
446
447
448
449
450
451
452
453
454
455
456
    Get the imports a given module makes.

    Args:
        module_fname (`str`):
            The name of the file of the module where we want to look at the imports (given relative to the root of
            the repo).
        cache (Dictionary `str` to `List[str]`, *optional*):
            To speed up this function if it was previously called on `module_fname`, the cache of all previously
            computed results.

    Returns:
        `List[str]`: The list of module filenames imported in the input `module_fname` (a submodule we import from that
        is a subfolder will give its init file).
457
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
458
459
460
461
    if cache is not None and module_fname in cache:
        return cache[module_fname]

    with open(PATH_TO_REPO / module_fname, "r", encoding="utf-8") as f:
462
463
        content = f.read()

464
465
    # Filter out all docstrings to not get imports in code examples. As before we need to deactivate formatting to
    # keep this as escaped quotes and avoid this function failing on this file.
466
467
468
    # fmt: off
    splits = content.split('\"\"\"')
    # fmt: on
Sylvain Gugger's avatar
Sylvain Gugger committed
469
470
471
    content = "".join(splits[::2])

    module_parts = str(module_fname).split(os.path.sep)
472
473
474
    imported_modules = []

    # Let's start with relative imports
Sylvain Gugger's avatar
Sylvain Gugger committed
475
476
477
478
479
480
481
    relative_imports = _re_single_line_relative_imports.findall(content)
    relative_imports = [
        (mod, imp) for mod, imp in relative_imports if "# tests_ignore" not in imp and imp.strip() != "("
    ]
    multiline_relative_imports = _re_multi_line_relative_imports.findall(content)
    relative_imports += [(mod, imp) for mod, imp in multiline_relative_imports if "# tests_ignore" not in imp]

482
    # We need to remove parts of the module name depending on the depth of the relative imports.
Sylvain Gugger's avatar
Sylvain Gugger committed
483
    for module, imports in relative_imports:
484
        level = 0
Sylvain Gugger's avatar
Sylvain Gugger committed
485
486
        while module.startswith("."):
            module = module[1:]
487
488
            level += 1

Sylvain Gugger's avatar
Sylvain Gugger committed
489
490
        if len(module) > 0:
            dep_parts = module_parts[: len(module_parts) - level] + module.split(".")
491
        else:
Sylvain Gugger's avatar
Sylvain Gugger committed
492
            dep_parts = module_parts[: len(module_parts) - level]
493
        imported_module = os.path.sep.join(dep_parts)
Sylvain Gugger's avatar
Sylvain Gugger committed
494
        imported_modules.append((imported_module, [imp.strip() for imp in imports.split(",")]))
495
496

    # Let's continue with direct imports
Sylvain Gugger's avatar
Sylvain Gugger committed
497
498
499
500
    direct_imports = _re_single_line_direct_imports.findall(content)
    direct_imports = [(mod, imp) for mod, imp in direct_imports if "# tests_ignore" not in imp and imp.strip() != "("]
    multiline_direct_imports = _re_multi_line_direct_imports.findall(content)
    direct_imports += [(mod, imp) for mod, imp in multiline_direct_imports if "# tests_ignore" not in imp]
501

502
    # We need to find the relative path of those imports.
Sylvain Gugger's avatar
Sylvain Gugger committed
503
    for module, imports in direct_imports:
504
        import_parts = module.split(".")[1:]  # ignore the name of the repo since we add it below.
Sylvain Gugger's avatar
Sylvain Gugger committed
505
506
507
        dep_parts = ["src", "transformers"] + import_parts
        imported_module = os.path.sep.join(dep_parts)
        imported_modules.append((imported_module, [imp.strip() for imp in imports.split(",")]))
508

Sylvain Gugger's avatar
Sylvain Gugger committed
509
    result = []
510
    # Double check we get proper modules (either a python file or a folder with an init).
Sylvain Gugger's avatar
Sylvain Gugger committed
511
512
513
514
515
516
517
518
    for module_file, imports in imported_modules:
        if (PATH_TO_REPO / f"{module_file}.py").is_file():
            module_file = f"{module_file}.py"
        elif (PATH_TO_REPO / module_file).is_dir() and (PATH_TO_REPO / module_file / "__init__.py").is_file():
            module_file = os.path.sep.join([module_file, "__init__.py"])
        imports = [imp for imp in imports if len(imp) > 0 and re.match("^[A-Za-z0-9_]*$", imp)]
        if len(imports) > 0:
            result.append((module_file, imports))
519

Sylvain Gugger's avatar
Sylvain Gugger committed
520
521
    if cache is not None:
        cache[module_fname] = result
Sylvain Gugger's avatar
Sylvain Gugger committed
522

Sylvain Gugger's avatar
Sylvain Gugger committed
523
    return result
524

525

526
def get_module_dependencies(module_fname: str, cache: Dict[str, List[str]] = None) -> List[str]:
Sylvain Gugger's avatar
Sylvain Gugger committed
527
    """
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
    Refines the result of `extract_imports` to remove subfolders and get a proper list of module filenames: if a file
    as an import `from utils import Foo, Bar`, with `utils` being a subfolder containing many files, this will traverse
    the `utils` init file to check where those dependencies come from: for instance the files utils/foo.py and utils/bar.py.

    Warning: This presupposes that all intermediate inits are properly built (with imports from the respective
    submodules) and work better if objects are defined in submodules and not the intermediate init (otherwise the
    intermediate init is added, and inits usually have a lot of dependencies).

    Args:
        module_fname (`str`):
            The name of the file of the module where we want to look at the imports (given relative to the root of
            the repo).
        cache (Dictionary `str` to `List[str]`, *optional*):
            To speed up this function if it was previously called on `module_fname`, the cache of all previously
            computed results.

    Returns:
        `List[str]`: The list of module filenames imported in the input `module_fname` (with submodule imports refined).
Sylvain Gugger's avatar
Sylvain Gugger committed
546
547
548
    """
    dependencies = []
    imported_modules = extract_imports(module_fname, cache=cache)
549
    # The while loop is to recursively traverse all inits we may encounter: we will add things as we go.
Sylvain Gugger's avatar
Sylvain Gugger committed
550
551
552
553
554
555
556
557
558
    while len(imported_modules) > 0:
        new_modules = []
        for module, imports in imported_modules:
            # If we end up in an __init__ we are often not actually importing from this init (except in the case where
            # the object is fully defined in the __init__)
            if module.endswith("__init__.py"):
                # So we get the imports from that init then try to find where our objects come from.
                new_imported_modules = extract_imports(module, cache=cache)
                for new_module, new_imports in new_imported_modules:
559
                    if any(i in new_imports for i in imports):
Sylvain Gugger's avatar
Sylvain Gugger committed
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
                        if new_module not in dependencies:
                            new_modules.append((new_module, [i for i in new_imports if i in imports]))
                        imports = [i for i in imports if i not in new_imports]
                if len(imports) > 0:
                    # If there are any objects lefts, they may be a submodule
                    path_to_module = PATH_TO_REPO / module.replace("__init__.py", "")
                    dependencies.extend(
                        [
                            os.path.join(module.replace("__init__.py", ""), f"{i}.py")
                            for i in imports
                            if (path_to_module / f"{i}.py").is_file()
                        ]
                    )
                    imports = [i for i in imports if not (path_to_module / f"{i}.py").is_file()]
                    if len(imports) > 0:
                        # Then if there are still objects left, they are fully defined in the init, so we keep it as a
                        # dependency.
                        dependencies.append(module)
            else:
                dependencies.append(module)
580

Sylvain Gugger's avatar
Sylvain Gugger committed
581
582
        imported_modules = new_modules
    return dependencies
Sylvain Gugger's avatar
Sylvain Gugger committed
583
584


585
def create_reverse_dependency_tree() -> List[Tuple[str, str]]:
586
587
588
    """
    Create a list of all edges (a, b) which mean that modifying a impacts b with a going over all module and test files.
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
589
590
591
592
    cache = {}
    all_modules = list(PATH_TO_TRANFORMERS.glob("**/*.py")) + list(PATH_TO_TESTS.glob("**/*.py"))
    all_modules = [str(mod.relative_to(PATH_TO_REPO)) for mod in all_modules]
    edges = [(dep, mod) for mod in all_modules for dep in get_module_dependencies(mod, cache=cache)]
593

Sylvain Gugger's avatar
Sylvain Gugger committed
594
    return list(set(edges))
595
596


597
def get_tree_starting_at(module: str, edges: List[Tuple[str, str]]) -> List[Union[str, List[str]]]:
598
    """
599
600
601
602
603
604
605
606
607
    Returns the tree starting at a given module following all edges.

    Args:
        module (`str`): The module that will be the root of the subtree we want.
        eges (`List[Tuple[str, str]]`): The list of all edges of the tree.

    Returns:
        `List[Union[str, List[str]]]`: The tree to print in the following format: [module, [list of edges
        starting at module], [list of edges starting at the preceding level], ...]
608
609
    """
    vertices_seen = [module]
Sylvain Gugger's avatar
Sylvain Gugger committed
610
    new_edges = [edge for edge in edges if edge[0] == module and edge[1] != module and "__init__.py" not in edge[1]]
611
612
613
    tree = [module]
    while len(new_edges) > 0:
        tree.append(new_edges)
614
        final_vertices = list({edge[1] for edge in new_edges})
615
        vertices_seen.extend(final_vertices)
Sylvain Gugger's avatar
Sylvain Gugger committed
616
617
618
619
620
        new_edges = [
            edge
            for edge in edges
            if edge[0] in final_vertices and edge[1] not in vertices_seen and "__init__.py" not in edge[1]
        ]
621
622
623
624
625
626
627

    return tree


def print_tree_deps_of(module, all_edges=None):
    """
    Prints the tree of modules depending on a given module.
628
629
630
631
632

    Args:
        module (`str`): The module that will be the root of the subtree we want.
        all_eges (`List[Tuple[str, str]]`, *optional*):
            The list of all edges of the tree. Will be set to `create_reverse_dependency_tree()` if not passed.
633
634
635
636
637
638
639
640
641
642
    """
    if all_edges is None:
        all_edges = create_reverse_dependency_tree()
    tree = get_tree_starting_at(module, all_edges)

    # The list of lines is a list of tuples (line_to_be_printed, module)
    # Keeping the modules lets us know where to insert each new lines in the list.
    lines = [(tree[0], tree[0])]
    for index in range(1, len(tree)):
        edges = tree[index]
643
        start_edges = {edge[0] for edge in edges}
644
645

        for start in start_edges:
646
            end_edges = {edge[1] for edge in edges if edge[0] == start}
647
648
649
650
651
652
653
654
655
656
657
            # We will insert all those edges just after the line showing start.
            pos = 0
            while lines[pos][1] != start:
                pos += 1
            lines = lines[: pos + 1] + [(" " * (2 * index) + end, end) for end in end_edges] + lines[pos + 1 :]

    for line in lines:
        # We don't print the refs that where just here to help build lines.
        print(line[0])


658
def init_test_examples_dependencies() -> Tuple[Dict[str, List[str]], List[str]]:
659
660
    """
    The test examples do not import from the examples (which are just scripts, not modules) so we need som extra
661
662
663
664
665
666
667
    care initializing the dependency map, which is the goal of this function. It initializes the dependency map for
    example files by linking each example to the example test file for the example framework.

    Returns:
        `Tuple[Dict[str, List[str]], List[str]]`: A tuple with two elements: the initialized dependency map which is a
        dict test example file to list of example files potentially tested by that test file, and the list of all
        example files (to avoid recomputing it later).
668
669
670
671
672
673
    """
    test_example_deps = {}
    all_examples = []
    for framework in ["flax", "pytorch", "tensorflow"]:
        test_files = list((PATH_TO_EXAMPLES / framework).glob("test_*.py"))
        all_examples.extend(test_files)
674
675
        # Remove the files at the root of examples/framework since they are not proper examples (they are eith utils
        # or example test files).
676
677
678
679
680
681
682
        examples = [
            f for f in (PATH_TO_EXAMPLES / framework).glob("**/*.py") if f.parent != PATH_TO_EXAMPLES / framework
        ]
        all_examples.extend(examples)
        for test_file in test_files:
            with open(test_file, "r", encoding="utf-8") as f:
                content = f.read()
683
            # Map all examples to the test files found in examples/framework.
684
685
686
            test_example_deps[str(test_file.relative_to(PATH_TO_REPO))] = [
                str(e.relative_to(PATH_TO_REPO)) for e in examples if e.name in content
            ]
687
            # Also map the test files to themselves.
688
689
690
            test_example_deps[str(test_file.relative_to(PATH_TO_REPO))].append(
                str(test_file.relative_to(PATH_TO_REPO))
            )
691
692
693
    return test_example_deps, all_examples


694
def create_reverse_dependency_map() -> Dict[str, List[str]]:
695
    """
696
697
698
699
700
701
    Create the dependency map from module/test filename to the list of modules/tests that depend on it recursively.

    Returns:
        `Dict[str, List[str]]`: The reverse dependency map as a dictionary mapping filenames to all the filenames
        depending on it recursively. This way the tests impacted by a change in file A are the test files in the list
        corresponding to key A in this result.
702
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
703
    cache = {}
704
    # Start from the example deps init.
705
    example_deps, examples = init_test_examples_dependencies()
706
    # Add all modules and all tests to all examples
707
    all_modules = list(PATH_TO_TRANFORMERS.glob("**/*.py")) + list(PATH_TO_TESTS.glob("**/*.py")) + examples
Sylvain Gugger's avatar
Sylvain Gugger committed
708
    all_modules = [str(mod.relative_to(PATH_TO_REPO)) for mod in all_modules]
709
    # Compute the direct dependencies of all modules.
Sylvain Gugger's avatar
Sylvain Gugger committed
710
    direct_deps = {m: get_module_dependencies(m, cache=cache) for m in all_modules}
711
    direct_deps.update(example_deps)
Sylvain Gugger's avatar
Sylvain Gugger committed
712

713
714
715
716
    # This recurses the dependencies
    something_changed = True
    while something_changed:
        something_changed = False
Sylvain Gugger's avatar
Sylvain Gugger committed
717
        for m in all_modules:
718
            for d in direct_deps[m]:
719
720
                # We stop recursing at an init (cause we always end up in the main init and we don't want to add all
                # files which the main init imports)
Sylvain Gugger's avatar
Sylvain Gugger committed
721
722
                if d.endswith("__init__.py"):
                    continue
723
724
                if d not in direct_deps:
                    raise ValueError(f"KeyError:{d}. From {m}")
Sylvain Gugger's avatar
Sylvain Gugger committed
725
726
727
728
                new_deps = set(direct_deps[d]) - set(direct_deps[m])
                if len(new_deps) > 0:
                    direct_deps[m].extend(list(new_deps))
                    something_changed = True
729
730
731

    # Finally we can build the reverse map.
    reverse_map = collections.defaultdict(list)
Sylvain Gugger's avatar
Sylvain Gugger committed
732
    for m in all_modules:
733
734
735
        for d in direct_deps[m]:
            reverse_map[d].append(m)

736
737
    # For inits, we don't do the reverse deps but the direct deps: if modifying an init, we want to make sure we test
    # all the modules impacted by that init.
Sylvain Gugger's avatar
Sylvain Gugger committed
738
739
740
741
742
    for m in [f for f in all_modules if f.endswith("__init__.py")]:
        direct_deps = get_module_dependencies(m, cache=cache)
        deps = sum([reverse_map[d] for d in direct_deps if not d.endswith("__init__.py")], direct_deps)
        reverse_map[m] = list(set(deps) - {m})

743
744
745
    return reverse_map


746
747
748
def create_module_to_test_map(
    reverse_map: Dict[str, List[str]] = None, filter_models: bool = False
) -> Dict[str, List[str]]:
Sylvain Gugger's avatar
Sylvain Gugger committed
749
750
    """
    Extract the tests from the reverse_dependency_map and potentially filters the model tests.
751
752
753
754
755
756
757
758
759
760

    Args:
        reverse_map (`Dict[str, List[str]]`, *optional*):
            The reverse dependency map as created by `create_reverse_dependency_map`. Will default to the result of
            that function if not provided.
        filter_models (`bool`, *optional*, defaults to `False`):
            Whether or not to filter model tests to only include core models if a file impacts a lot of models.

    Returns:
        `Dict[str, List[str]]`: A dictionary that maps each file to the tests to execute if that file was modified.
Sylvain Gugger's avatar
Sylvain Gugger committed
761
762
763
    """
    if reverse_map is None:
        reverse_map = create_reverse_dependency_map()
764

765
    # Utility that tells us if a given file is a test (taking test examples into account)
766
767
768
769
770
771
772
    def is_test(fname):
        if fname.startswith("tests"):
            return True
        if fname.startswith("examples") and fname.split(os.path.sep)[-1].startswith("test"):
            return True
        return False

773
    # Build the test map
774
    test_map = {module: [f for f in deps if is_test(f)] for module, deps in reverse_map.items()}
775

Sylvain Gugger's avatar
Sylvain Gugger committed
776
777
    if not filter_models:
        return test_map
778

779
    # Now we deal with the filtering if `filter_models` is True.
Sylvain Gugger's avatar
Sylvain Gugger committed
780
781
782
    num_model_tests = len(list(PATH_TO_TESTS.glob("models/*")))

    def has_many_models(tests):
783
        # We filter to core models when a given file impacts more than half the model tests.
Sylvain Gugger's avatar
Sylvain Gugger committed
784
785
786
787
788
789
790
        model_tests = {Path(t).parts[2] for t in tests if t.startswith("tests/models/")}
        return len(model_tests) > num_model_tests // 2

    def filter_tests(tests):
        return [t for t in tests if not t.startswith("tests/models/") or Path(t).parts[2] in IMPORTANT_MODELS]

    return {module: (filter_tests(tests) if has_many_models(tests) else tests) for module, tests in test_map.items()}
791
792


Sylvain Gugger's avatar
Sylvain Gugger committed
793
def check_imports_all_exist():
794
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
795
    Isn't used per se by the test fetcher but might be used later as a quality check. Putting this here for now so the
796
    code is not lost. This checks all imports in a given file do exist.
797
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
798
799
800
801
    cache = {}
    all_modules = list(PATH_TO_TRANFORMERS.glob("**/*.py")) + list(PATH_TO_TESTS.glob("**/*.py"))
    all_modules = [str(mod.relative_to(PATH_TO_REPO)) for mod in all_modules]
    direct_deps = {m: get_module_dependencies(m, cache=cache) for m in all_modules}
802

Sylvain Gugger's avatar
Sylvain Gugger committed
803
804
805
806
807
808
    for module, deps in direct_deps.items():
        for dep in deps:
            if not (PATH_TO_REPO / dep).is_file():
                print(f"{module} has dependency on {dep} which does not exist.")


809
810
811
812
def _print_list(l) -> str:
    """
    Pretty print a list of elements with one line per element and a - starting each line.
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
813
814
815
    return "\n".join([f"- {f}" for f in l])


816
817
818
819
820
821
822
823
def create_json_map(test_files_to_run: List[str], json_output_file: str):
    """
    Creates a map from a list of tests to run to easily split them by category, when running parallelism of slow tests.

    Args:
        test_files_to_run (`List[str]`): The list of tests to run.
        json_output_file (`str`): The path where to store the built json map.
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
    if json_output_file is None:
        return

    test_map = {}
    for test_file in test_files_to_run:
        # `test_file` is a path to a test folder/file, starting with `tests/`. For example,
        #   - `tests/models/bert/test_modeling_bert.py` or `tests/models/bert`
        #   - `tests/trainer/test_trainer.py` or `tests/trainer`
        #   - `tests/test_modeling_common.py`
        names = test_file.split(os.path.sep)
        if names[1] == "models":
            # take the part like `models/bert` for modeling tests
            key = os.path.sep.join(names[1:3])
        elif len(names) > 2 or not test_file.endswith(".py"):
            # test folders under `tests` or python files under them
            # take the part like tokenization, `pipeline`, etc. for other test categories
            key = os.path.sep.join(names[1:2])
841
        else:
Sylvain Gugger's avatar
Sylvain Gugger committed
842
843
            # common test files directly under `tests/`
            key = "common"
844

Sylvain Gugger's avatar
Sylvain Gugger committed
845
846
847
        if key not in test_map:
            test_map[key] = []
        test_map[key].append(test_file)
848

Sylvain Gugger's avatar
Sylvain Gugger committed
849
850
851
852
853
    # sort the keys & values
    keys = sorted(test_map.keys())
    test_map = {k: " ".join(sorted(test_map[k])) for k in keys}
    with open(json_output_file, "w", encoding="UTF-8") as fp:
        json.dump(test_map, fp, ensure_ascii=False)
854
855


856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
def infer_tests_to_run(
    output_file: str,
    diff_with_last_commit: bool = False,
    filter_models: bool = True,
    json_output_file: Optional[str] = None,
):
    """
    The main function called by the test fetcher. Determines the tests to run from the diff.

    Args:
        output_file (`str`):
            The path where to store the summary of the test fetcher analysis. Other files will be stored in the same
            folder:

            - examples_test_list.txt: The list of examples tests to run.
            - test_repo_utils.txt: Will indicate if the repo utils tests should be run or not.
            - doctest_list.txt: The list of doctests to run.

        diff_with_last_commit (`bool`, *optional*, defaults to `False`):
            Whether to analyze the diff with the last commit (for use on the main branch after a PR is merged) or with
            the branching point from main (for use on each PR).
        filter_models (`bool`, *optional*, defaults to `True`):
            Whether or not to filter the tests to core models only, when a file modified results in a lot of model
            tests.
        json_output_file (`str`, *optional*):
            The path where to store the json file mapping categories of tests to tests to run (used for parallelism or
            the slow tests).
    """
884
    modified_files = get_modified_python_files(diff_with_last_commit=diff_with_last_commit)
885
886
887
    print(f"\n### MODIFIED FILES ###\n{_print_list(modified_files)}")

    # Create the map that will give us all impacted modules.
Sylvain Gugger's avatar
Sylvain Gugger committed
888
    reverse_map = create_reverse_dependency_map()
889
890
    impacted_files = modified_files.copy()
    for f in modified_files:
Sylvain Gugger's avatar
Sylvain Gugger committed
891
892
        if f in reverse_map:
            impacted_files.extend(reverse_map[f])
893
894

    # Remove duplicates
895
    impacted_files = sorted(set(impacted_files))
896
897
898
    print(f"\n### IMPACTED FILES ###\n{_print_list(impacted_files)}")

    # Grab the corresponding test files:
Sylvain Gugger's avatar
Sylvain Gugger committed
899
    if "setup.py" in modified_files:
900
        test_files_to_run = ["tests", "examples"]
Sylvain Gugger's avatar
Sylvain Gugger committed
901
        repo_utils_launch = True
902
    else:
Sylvain Gugger's avatar
Sylvain Gugger committed
903
904
905
906
907
908
909
910
911
        # All modified tests need to be run.
        test_files_to_run = [
            f for f in modified_files if f.startswith("tests") and f.split(os.path.sep)[-1].startswith("test")
        ]
        # Then we grab the corresponding test files.
        test_map = create_module_to_test_map(reverse_map=reverse_map, filter_models=filter_models)
        for f in modified_files:
            if f in test_map:
                test_files_to_run.extend(test_map[f])
912
        test_files_to_run = sorted(set(test_files_to_run))
913
914
        # Remove repo utils tests
        test_files_to_run = [f for f in test_files_to_run if not f.split(os.path.sep)[1] == "repo_utils"]
Sylvain Gugger's avatar
Sylvain Gugger committed
915
916
        # Remove SageMaker tests
        test_files_to_run = [f for f in test_files_to_run if not f.split(os.path.sep)[1] == "sagemaker"]
917
        # Make sure we did not end up with a test file that was removed
Sylvain Gugger's avatar
Sylvain Gugger committed
918
919
        test_files_to_run = [f for f in test_files_to_run if (PATH_TO_REPO / f).exists()]

920
        repo_utils_launch = any(f.split(os.path.sep)[0] == "utils" for f in modified_files)
Sylvain Gugger's avatar
Sylvain Gugger committed
921
922
923
924
925

    if repo_utils_launch:
        repo_util_file = Path(output_file).parent / "test_repo_utils.txt"
        with open(repo_util_file, "w", encoding="utf-8") as f:
            f.write("tests/repo_utils")
926

927
928
    examples_tests_to_run = [f for f in test_files_to_run if f.startswith("examples")]
    test_files_to_run = [f for f in test_files_to_run if not f.startswith("examples")]
929
930
931
932
933
    print(f"\n### TEST TO RUN ###\n{_print_list(test_files_to_run)}")
    if len(test_files_to_run) > 0:
        with open(output_file, "w", encoding="utf-8") as f:
            f.write(" ".join(test_files_to_run))

Yih-Dar's avatar
Yih-Dar committed
934
935
936
937
938
939
940
        # Create a map that maps test categories to test files, i.e. `models/bert` -> [...test_modeling_bert.py, ...]

        # Get all test directories (and some common test files) under `tests` and `tests/models` if `test_files_to_run`
        # contains `tests` (i.e. when `setup.py` is changed).
        if "tests" in test_files_to_run:
            test_files_to_run = get_all_tests()

Sylvain Gugger's avatar
Sylvain Gugger committed
941
        create_json_map(test_files_to_run, json_output_file)
Yih-Dar's avatar
Yih-Dar committed
942

943
944
    print(f"\n### EXAMPLES TEST TO RUN ###\n{_print_list(examples_tests_to_run)}")
    if len(examples_tests_to_run) > 0:
Yih-Dar's avatar
Yih-Dar committed
945
946
947
        # We use `all` in the case `commit_flags["test_all"]` as well as in `create_circleci_config.py` for processing
        if examples_tests_to_run == ["examples"]:
            examples_tests_to_run = ["all"]
948
949
950
951
        example_file = Path(output_file).parent / "examples_test_list.txt"
        with open(example_file, "w", encoding="utf-8") as f:
            f.write(" ".join(examples_tests_to_run))

952
953
954
955
956
957
958
959
    doctest_list = get_doctest_files()

    print(f"\n### DOCTEST TO RUN ###\n{_print_list(doctest_list)}")
    if len(doctest_list) > 0:
        doctest_file = Path(output_file).parent / "doctest_list.txt"
        with open(doctest_file, "w", encoding="utf-8") as f:
            f.write(" ".join(doctest_list))

960

961
def filter_tests(output_file: str, filters: List[str]):
Sylvain Gugger's avatar
Sylvain Gugger committed
962
963
964
965
966
967
968
    """
    Reads the content of the output file and filters out all the tests in a list of given folders.

    Args:
        output_file (`str` or `os.PathLike`): The path to the output file of the tests fetcher.
        filters (`List[str]`): A list of folders to filter.
    """
969
970
971
972
973
974
    if not os.path.isfile(output_file):
        print("No test file found.")
        return
    with open(output_file, "r", encoding="utf-8") as f:
        test_files = f.read().split(" ")

Sylvain Gugger's avatar
Sylvain Gugger committed
975
    if len(test_files) == 0 or test_files == [""]:
976
977
        print("No tests to filter.")
        return
Sylvain Gugger's avatar
Sylvain Gugger committed
978

979
    if test_files == ["tests"]:
Sylvain Gugger's avatar
Sylvain Gugger committed
980
        test_files = [os.path.join("tests", f) for f in os.listdir("tests") if f not in ["__init__.py"] + filters]
981
    else:
Sylvain Gugger's avatar
Sylvain Gugger committed
982
        test_files = [f for f in test_files if f.split(os.path.sep)[1] not in filters]
983
984
985
986
987

    with open(output_file, "w", encoding="utf-8") as f:
        f.write(" ".join(test_files))


988
def parse_commit_message(commit_message: str) -> Dict[str, bool]:
Sylvain Gugger's avatar
Sylvain Gugger committed
989
990
991
    """
    Parses the commit message to detect if a command is there to skip, force all or part of the CI.

992
993
994
995
996
997
    Args:
        commit_message (`str`): The commit message of the current commit.

    Returns:
        `Dict[str, bool]`: A dictionary of strings to bools with keys the following keys: `"skip"`,
        `"test_all_models"` and `"test_all"`.
Sylvain Gugger's avatar
Sylvain Gugger committed
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
    """
    if commit_message is None:
        return {"skip": False, "no_filter": False, "test_all": False}

    command_search = re.search(r"\[([^\]]*)\]", commit_message)
    if command_search is not None:
        command = command_search.groups()[0]
        command = command.lower().replace("-", " ").replace("_", " ")
        skip = command in ["ci skip", "skip ci", "circleci skip", "skip circleci"]
        no_filter = set(command.split(" ")) == {"no", "filter"}
        test_all = set(command.split(" ")) == {"test", "all"}
        return {"skip": skip, "no_filter": no_filter, "test_all": test_all}
    else:
        return {"skip": False, "no_filter": False, "test_all": False}


1014
1015
1016
1017
1018
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--output_file", type=str, default="test_list.txt", help="Where to store the list of tests to run"
    )
Yih-Dar's avatar
Yih-Dar committed
1019
1020
1021
1022
1023
1024
    parser.add_argument(
        "--json_output_file",
        type=str,
        default="test_map.json",
        help="Where to store the tests to run in a dictionary format mapping test categories to test files",
    )
1025
1026
1027
1028
1029
    parser.add_argument(
        "--diff_with_last_commit",
        action="store_true",
        help="To fetch the tests between the current commit and the last commit",
    )
1030
    parser.add_argument(
Sylvain Gugger's avatar
Sylvain Gugger committed
1031
        "--filter_tests",
1032
        action="store_true",
Sylvain Gugger's avatar
Sylvain Gugger committed
1033
        help="Will filter the pipeline/repo utils tests outside of the generated list of tests.",
1034
    )
1035
1036
1037
1038
1039
1040
    parser.add_argument(
        "--print_dependencies_of",
        type=str,
        help="Will only print the tree of modules depending on the file passed.",
        default=None,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
1041
1042
1043
1044
1045
1046
    parser.add_argument(
        "--commit_message",
        type=str,
        help="The commit message (which could contain a command to force all tests or skip the CI).",
        default=None,
    )
1047
    args = parser.parse_args()
1048
1049
    if args.print_dependencies_of is not None:
        print_tree_deps_of(args.print_dependencies_of)
Sylvain Gugger's avatar
Sylvain Gugger committed
1050
1051
    elif args.filter_tests:
        filter_tests(args.output_file, ["pipelines", "repo_utils"])
1052
    else:
Sylvain Gugger's avatar
Sylvain Gugger committed
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
        repo = Repo(PATH_TO_REPO)
        commit_message = repo.head.commit.message
        commit_flags = parse_commit_message(commit_message)
        if commit_flags["skip"]:
            print("Force-skipping the CI")
            quit()
        if commit_flags["no_filter"]:
            print("Running all tests fetched without filtering.")
        if commit_flags["test_all"]:
            print("Force-launching all tests")
1063
1064

        diff_with_last_commit = args.diff_with_last_commit
1065
1066
        if not diff_with_last_commit and not repo.head.is_detached and repo.head.ref == repo.refs.main:
            print("main branch detected, fetching tests against last commit.")
1067
1068
            diff_with_last_commit = True

Sylvain Gugger's avatar
Sylvain Gugger committed
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
        if not commit_flags["test_all"]:
            try:
                infer_tests_to_run(
                    args.output_file,
                    diff_with_last_commit=diff_with_last_commit,
                    json_output_file=args.json_output_file,
                    filter_models=not commit_flags["no_filter"],
                )
                filter_tests(args.output_file, ["repo_utils"])
            except Exception as e:
                print(f"\nError when trying to grab the relevant tests: {e}\n\nRunning all tests.")
                commit_flags["test_all"] = True

        if commit_flags["test_all"]:
1083
            with open(args.output_file, "w", encoding="utf-8") as f:
Sylvain Gugger's avatar
Sylvain Gugger committed
1084
1085
1086
1087
                f.write("tests")
            example_file = Path(args.output_file).parent / "examples_test_list.txt"
            with open(example_file, "w", encoding="utf-8") as f:
                f.write("all")
Sylvain Gugger's avatar
Sylvain Gugger committed
1088
1089
1090

            test_files_to_run = get_all_tests()
            create_json_map(test_files_to_run, args.json_output_file)