pytorch-lightning

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

15
import os
16
import shutil
17
import tarfile
18

19

20
def clean_tarfile(file_path: str, mode: str) -> None:
21
    """This utility removes all files extracted from a tarfile."""
22
    if not os.path.exists(file_path):
23
        return
24

25
    with tarfile.open(file_path, mode=mode) as tar_ref:
26
        for member in tar_ref.getmembers():
27
            p = member.path
28
            if p == "." or not os.path.exists(p):
29
                continue
30
            try:
31
                if os.path.isfile(p):
32
                    os.remove(p)
33
                else:
34
                    shutil.rmtree(p)
35
            except (FileNotFoundError, OSError, PermissionError):
36
                pass
37

38
    if os.path.exists(file_path):
39
        os.remove(file_path)
40

41

42
def extract_tarfile(file_path: str, extract_path: str, mode: str) -> None:
43
    """This utility extracts all files from a tarfile."""
44
    if not os.path.exists(file_path):
45
        return
46

47
    with tarfile.open(file_path, mode=mode) as tar_ref:
48
        for member in tar_ref.getmembers():
49
            try:
50
                tar_ref.extract(member, path=extract_path, set_attrs=False)
51
            except PermissionError:
52
                raise PermissionError(f"Could not extract tar file {file_path}")
53