utils.py 1.54 KB
Newer Older
Johannes Bechberger's avatar
Johannes Bechberger committed
1 2 3
import logging
from os import path
import sys
4 5
from typing import Tuple, Optional
import re
6

7
COLOR_OUTPUT_IF_POSSIBLE = False
8

9 10 11 12 13 14
if sys.stdout.isatty():
    try:
        import termcolor
        COLOR_OUTPUT_IF_POSSIBLE = True
    except ImportError:
        COLOR_OUTPUT_IF_POSSIBLE = False
15 16


17
def force_colored_output():
18
    global COLOR_OUTPUT_IF_POSSIBLE
19
    COLOR_OUTPUT_IF_POSSIBLE = True
20

Johannes Bechberger's avatar
Johannes Bechberger committed
21 22

def get_mjtest_basedir() -> str:
23
    return path.dirname(path.dirname(path.dirname(path.realpath(__file__))))
Johannes Bechberger's avatar
Johannes Bechberger committed
24 25


26 27 28 29
def colored(text: str, *args, **kwargs):
    """
    Wrapper around termcolor.colored (if it's loadable)
    """
30 31 32 33
    global COLOR_OUTPUT_IF_POSSIBLE
    if COLOR_OUTPUT_IF_POSSIBLE:
        return termcolor.colored(text, *args, **kwargs)
    else:
34 35 36 37 38 39 40
        return text


def cprint(text: str, *args, **kwargs):
    """
    Wrapper around termcolor.cprint (if it's loadable)
    """
41 42 43 44
    global COLOR_OUTPUT_IF_POSSIBLE
    if COLOR_OUTPUT_IF_POSSIBLE:
        termcolor.cprint(text, *args, **kwargs)
    else:
45 46 47 48 49 50 51 52 53 54 55 56 57 58 59
        print(text)


def get_main_class_name(file: str) -> Optional[str]:
    current_class = None
    with open(file, "r") as f:
        for line in f:
            if line.startswith("class "):
                match = re.search("[A-Za-z_0-9]+", line.replace("class ", ""))
                if match:
                    has_public_class = True
                    current_class = match.group(0)
            elif "String[]" in line and "main" in line and "void" in line and "static" in line and "public" in line:
                return current_class
    return None