stable-diffusion-webui

1
import os
2
import re
3

4
from modules import shared
5
from modules.paths_internal import script_path, cwd
6

7

8
def natural_sort_key(s, regex=re.compile('([0-9]+)')):
9
    return [int(text) if text.isdigit() else text.lower() for text in regex.split(s)]
10

11

12
def listfiles(dirname):
13
    filenames = [os.path.join(dirname, x) for x in sorted(os.listdir(dirname), key=natural_sort_key) if not x.startswith(".")]
14
    return [file for file in filenames if os.path.isfile(file)]
15

16

17
def html_path(filename):
18
    return os.path.join(script_path, "html", filename)
19

20

21
def html(filename):
22
    path = html_path(filename)
23

24
    try:
25
        with open(path, encoding="utf8") as file:
26
            return file.read()
27
    except OSError:
28
        return ""
29

30

31
def walk_files(path, allowed_extensions=None):
32
    if not os.path.exists(path):
33
        return
34

35
    if allowed_extensions is not None:
36
        allowed_extensions = set(allowed_extensions)
37

38
    items = list(os.walk(path, followlinks=True))
39
    items = sorted(items, key=lambda x: natural_sort_key(x[0]))
40

41
    for root, _, files in items:
42
        for filename in sorted(files, key=natural_sort_key):
43
            if allowed_extensions is not None:
44
                _, ext = os.path.splitext(filename)
45
                if ext.lower() not in allowed_extensions:
46
                    continue
47

48
            if not shared.opts.list_hidden_files and ("/." in root or "\\." in root):
49
                continue
50

51
            yield os.path.join(root, filename)
52

53

54
def ldm_print(*args, **kwargs):
55
    if shared.opts.hide_ldm_prints:
56
        return
57

58
    print(*args, **kwargs)
59

60

61
def truncate_path(target_path, base_path=cwd):
62
    abs_target, abs_base = os.path.abspath(target_path), os.path.abspath(base_path)
63
    try:
64
        if os.path.commonpath([abs_target, abs_base]) == abs_base:
65
            return os.path.relpath(abs_target, abs_base)
66
    except ValueError:
67
        pass
68
    return abs_target
69

70

71
class MassFileListerCachedDir:
72
    """A class that caches file metadata for a specific directory."""
73

74
    def __init__(self, dirname):
75
        self.files = None
76
        self.files_cased = None
77
        self.dirname = dirname
78

79
        stats = ((x.name, x.stat(follow_symlinks=False)) for x in os.scandir(self.dirname))
80
        files = [(n, s.st_mtime, s.st_ctime) for n, s in stats]
81
        self.files = {x[0].lower(): x for x in files}
82
        self.files_cased = {x[0]: x for x in files}
83

84

85
class MassFileLister:
86
    """A class that provides a way to check for the existence and mtime/ctile of files without doing more than one stat call per file."""
87

88
    def __init__(self):
89
        self.cached_dirs = {}
90

91
    def find(self, path):
92
        """
93
        Find the metadata for a file at the given path.
94

95
        Returns:
96
            tuple or None: A tuple of (name, mtime, ctime) if the file exists, or None if it does not.
97
        """
98

99
        dirname, filename = os.path.split(path)
100

101
        cached_dir = self.cached_dirs.get(dirname)
102
        if cached_dir is None:
103
            cached_dir = MassFileListerCachedDir(dirname)
104
            self.cached_dirs[dirname] = cached_dir
105

106
        stats = cached_dir.files_cased.get(filename)
107
        if stats is not None:
108
            return stats
109

110
        stats = cached_dir.files.get(filename.lower())
111
        if stats is None:
112
            return None
113

114
        try:
115
            os_stats = os.stat(path, follow_symlinks=False)
116
            return filename, os_stats.st_mtime, os_stats.st_ctime
117
        except Exception:
118
            return None
119

120
    def exists(self, path):
121
        """Check if a file exists at the given path."""
122

123
        return self.find(path) is not None
124

125
    def mctime(self, path):
126
        """
127
        Get the modification and creation times for a file at the given path.
128

129
        Returns:
130
            tuple: A tuple of (mtime, ctime) if the file exists, or (0, 0) if it does not.
131
        """
132

133
        stats = self.find(path)
134
        return (0, 0) if stats is None else stats[1:3]
135

136
    def reset(self):
137
        """Clear the cache of all directories."""
138
        self.cached_dirs.clear()
139