ClickHouse/docs/tools/util.py
2022-08-08 01:06:47 +02:00

137 lines
3.5 KiB
Python

import collections
import contextlib
import datetime
import multiprocessing
import os
import shutil
import sys
import socket
import tempfile
import threading
import jinja2
import yaml
@contextlib.contextmanager
def temp_dir():
path = tempfile.mkdtemp(dir=os.environ.get("TEMP"))
try:
yield path
finally:
shutil.rmtree(path)
@contextlib.contextmanager
def cd(new_cwd):
old_cwd = os.getcwd()
os.chdir(new_cwd)
try:
yield
finally:
os.chdir(old_cwd)
def get_free_port():
with contextlib.closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
s.bind(("", 0))
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
return s.getsockname()[1]
def run_function_in_parallel(func, args_list, threads=False):
processes = []
exit_code = 0
for task in args_list:
cls = threading.Thread if threads else multiprocessing.Process
processes.append(cls(target=func, args=task))
processes[-1].start()
for process in processes:
process.join()
if not threads:
if process.exitcode and not exit_code:
exit_code = process.exitcode
if exit_code:
sys.exit(exit_code)
def read_md_file(path):
in_meta = False
meta = {}
meta_text = []
content = []
if os.path.exists(path):
with open(path, "r") as f:
for line in f:
if line.startswith("---"):
if in_meta:
in_meta = False
meta = yaml.full_load("".join(meta_text))
else:
in_meta = True
else:
if in_meta:
meta_text.append(line)
else:
content.append(line)
return meta, "".join(content)
def write_md_file(path, meta, content):
dirname = os.path.dirname(path)
if not os.path.exists(dirname):
os.makedirs(dirname)
with open(path, "w") as f:
if meta:
print("---", file=f)
yaml.dump(meta, f)
print("---", file=f)
if not content.startswith("\n"):
print("", file=f)
f.write(content)
def represent_ordereddict(dumper, data):
value = []
for item_key, item_value in data.items():
node_key = dumper.represent_data(item_key)
node_value = dumper.represent_data(item_value)
value.append((node_key, node_value))
return yaml.nodes.MappingNode("tag:yaml.org,2002:map", value)
yaml.add_representer(collections.OrderedDict, represent_ordereddict)
def init_jinja2_filters(env):
import website
chunk_size = 10240
env.filters["chunks"] = lambda line: [
line[i : i + chunk_size] for i in range(0, len(line), chunk_size)
]
env.filters["to_rfc882"] = lambda d: datetime.datetime.strptime(
d, "%Y-%m-%d"
).strftime("%a, %d %b %Y %H:%M:%S GMT")
def init_jinja2_env(args):
import mdx_clickhouse
env = jinja2.Environment(
loader=jinja2.FileSystemLoader(
[args.website_dir, os.path.join(args.src_dir, "docs", "_includes")]
),
extensions=["jinja2.ext.i18n", "jinja2_highlight.HighlightExtension"],
)
env.extend(jinja2_highlight_cssclass="syntax p-3 my-3")
translations_dir = os.path.join(args.website_dir, "locale")
env.install_gettext_translations(
mdx_clickhouse.get_translations(translations_dir, "en"), newstyle=True
)
init_jinja2_filters(env)
return env