<commit_before>#!/usr/bin/env python3
import json
import os
import sys
from getpass import getpass
from sys import argv
import getopt

from prettytable import PrettyTable

from . import colors
from . import config
from . import global_vars
from . import json_encoders
from . import utils
from .exceptions import PolygonNotLoginnedError
from .local_file import LocalFile
from .problem import ProblemSession


def fatal(error):
    print(error)
    exit(0)


def load_session():
    try:
        if os.path.exists(config.get_session_file_path()):
            session_data_json = open(config.get_session_file_path(), 'r').read()
        elif os.path.exists(os.path.join('..', config.get_session_file_path())):
            config.internal_directory_path = os.path.join('..', config.internal_directory_path)
            session_data_json = open(config.get_session_file_path(), 'r').read()
        else:
            return False
        session_data = json.loads(session_data_json, object_hook=json_encoders.my_json_decoder)
        global_vars.problem = ProblemSession(config.polygon_url, session_data["problemId"])
        global_vars.problem.use_ready_session(session_data)
        return True
    except:
        return False


def save_session():
    session_data = global_vars.problem.dump_session()
    session_data_json = json.dumps(session_data, sort_keys=True, indent='  ', default=json_encoders.my_json_encoder)
    utils.safe_rewrite_file(config.get_session_file_path(), session_data_json)


def process_init(args):
    problem_id = None
    try:
        problem_id = int(args[0])
    except:
        print_help()
    if config.login:
        print('Using login %s from config' % config.login)
    else:
        print('Enter login:', end=' ')
        sys.stdout.flush()
        config.login = sys.stdin.readline().strip()
    if config.password:
        print('Using password from config')
    else:
        config.password = getpass('Enter password: ')
    global_vars.problem = ProblemSession(config.polygon_url, problem_id)
    global_vars.problem.create_new_session(config.login, config.password)
    save_session()


def process_relogin(args):
    if len(args) != 0:
        print_help()
    if not load_session() or global_vars.problem.problem_id is None:
        fatal('No problemId known. Use init instead.')
    local_files = global_vars.problem.local_files
    process_init([global_vars.problem.problem_id])
    global_vars.problem.local_files = local_files
    save_session()


def process_update(args):
    if not load_session() or global_vars.problem.sessionId is None:
        fatal('No session known. Use relogin or init first.')
    files = global_vars.problem.get_all_files_list()
    table = PrettyTable(['File type', 'Polygon name', 'Local path', 'Status'])
    flat = False
    if args and args[0] == '--flat':  # TODO: getopt or something like it
        flat = True
        args = args[1:]
    for file in files:
        if file.type == 'resource':
            continue
        local_file = global_vars.problem.get_local_by_polygon(file)
        need_file = file.name in args or \
                    local_file is not None and \
                    (local_file.name in args or local_file.filename in args or local_file.get_path() in args)
        if args and not need_file:
            continue
        if local_file is not None:
            print('Updating local file %s from %s' % (local_file.name, file.name))
            status = utils.safe_update_file(local_file.get_internal_path(),
                                            local_file.get_path(),
                                            file.get_content()
                                            )
            if status == 'Not changed':
                status = colors.info(status)
            elif status == 'Conflict':
                status = colors.error(status)
            else:
                status = colors.warning(status)
        else:
            status = colors.success('New')
            local_file = LocalFile()
            local_file.name = file.name.split('.')[0]
            local_file.dir = '' if flat else file.get_default_local_dir()
            local_file.type = file.type
            local_file.filename = file.name
            local_file.polygon_filename = file.name
            print('Downloading new file %s to %s' % (file.name, local_file.get_path()))
            content = file.get_content()
            utils.safe_rewrite_file(local_file.get_path(), content)
            utils.safe_rewrite_file(local_file.get_internal_path(), content)
            global_vars.problem.local_files.append(local_file)
        table.add_row([file.type, file.name, local_file.get_path(), status])
    print(table)
    save_session()


def process_add(args):
    if not load_session() or global_vars.problem.sessionId is None:
        fatal('No session known. Use relogin or init first.')
    if len(args) < 2:
        print_help()
    valid_types = ['solution', 'source', 'checker', 'validator']
    try:
        options, args = getopt.gnu_getopt(args, "t:")
    except getopt.GetoptError as err:
        fatal(err.msg)
    if args[0] not in valid_types:
        fatal('type should be in ' + str(valid_types))
    as_checker = False
    as_validator = False
    solution_type = None
    for opt in options:
        if opt[0] == '-t':
            solution_type = opt[1]
        else:
            pass # will be used with more options
    if args[0] == 'checker' or args[0] == 'validator':
        if len(args) != 2:
            fatal('can''t set several ' + args[0] + 's')
        if args[0] == 'checker':
            as_checker = True
        else:
            as_validator = True
        args[0] = 'source'
    if (solution_type is not None) and (args[0] != 'solution'):
        fatal('solution type can be set only on solutions')
    table = PrettyTable(['File type', 'Polygon name', 'Local path', 'Status'])
    for filename in args[1:]:
        local = global_vars.problem.get_local_by_filename(os.path.basename(filename))
        if local is not None:
            print('file %s already added, use commit instead' % os.path.basename(filename))
            status = colors.warning('Already added')
        else:
            local = LocalFile(os.path.basename(filename),
                              os.path.dirname(filename),
                              str(os.path.basename(filename).split('.')[0]),
                              args[0]
                              )
            if local.upload():
                status = colors.success('Uploaded')
                global_vars.problem.local_files.append(local)
                if as_checker:
                    global_vars.problem.set_checker_validator(local.polygon_filename, 'checker')
                if as_validator:
                    global_vars.problem.set_checker_validator(local.polygon_filename, 'validator')
                if solution_type is not None:
                    valid_solution_types = ['MA', 'OK', 'RJ', 'TL', 'WA', 'PE', 'ML', 'RE']
                    if solution_type.upper() in valid_solution_types:
                        global_vars.problem.change_solution_type(local.polygon_filename, solution_type.upper())
                    else:
                        status = colors.warning('Uploaded')
                        print('solytion type should be in ' + str(valid_solution_types))
            else:<commit_msg>support MAIN solution type (alias to MA) and change error message<commit_after>#!/usr/bin/env python3
import json
import os
import sys
from getpass import getpass
from sys import argv
import getopt

from prettytable import PrettyTable

from . import colors
from . import config
from . import global_vars
from . import json_encoders
from . import utils
from .exceptions import PolygonNotLoginnedError
from .local_file import LocalFile
from .problem import ProblemSession


def fatal(error):
    print(error)
    exit(0)


def load_session():
    try:
        if os.path.exists(config.get_session_file_path()):
            session_data_json = open(config.get_session_file_path(), 'r').read()
        elif os.path.exists(os.path.join('..', config.get_session_file_path())):
            config.internal_directory_path = os.path.join('..', config.internal_directory_path)
            session_data_json = open(config.get_session_file_path(), 'r').read()
        else:
            return False
        session_data = json.loads(session_data_json, object_hook=json_encoders.my_json_decoder)
        global_vars.problem = ProblemSession(config.polygon_url, session_data["problemId"])
        global_vars.problem.use_ready_session(session_data)
        return True
    except:
        return False


def save_session():
    session_data = global_vars.problem.dump_session()
    session_data_json = json.dumps(session_data, sort_keys=True, indent='  ', default=json_encoders.my_json_encoder)
    utils.safe_rewrite_file(config.get_session_file_path(), session_data_json)


def process_init(args):
    problem_id = None
    try:
        problem_id = int(args[0])
    except:
        print_help()
    if config.login:
        print('Using login %s from config' % config.login)
    else:
        print('Enter login:', end=' ')
        sys.stdout.flush()
        config.login = sys.stdin.readline().strip()
    if config.password:
        print('Using password from config')
    else:
        config.password = getpass('Enter password: ')
    global_vars.problem = ProblemSession(config.polygon_url, problem_id)
    global_vars.problem.create_new_session(config.login, config.password)
    save_session()


def process_relogin(args):
    if len(args) != 0:
        print_help()
    if not load_session() or global_vars.problem.problem_id is None:
        fatal('No problemId known. Use init instead.')
    local_files = global_vars.problem.local_files
    process_init([global_vars.problem.problem_id])
    global_vars.problem.local_files = local_files
    save_session()


def process_update(args):
    if not load_session() or global_vars.problem.sessionId is None:
        fatal('No session known. Use relogin or init first.')
    files = global_vars.problem.get_all_files_list()
    table = PrettyTable(['File type', 'Polygon name', 'Local path', 'Status'])
    flat = False
    if args and args[0] == '--flat':  # TODO: getopt or something like it
        flat = True
        args = args[1:]
    for file in files:
        if file.type == 'resource':
            continue
        local_file = global_vars.problem.get_local_by_polygon(file)
        need_file = file.name in args or \
                    local_file is not None and \
                    (local_file.name in args or local_file.filename in args or local_file.get_path() in args)
        if args and not need_file:
            continue
        if local_file is not None:
            print('Updating local file %s from %s' % (local_file.name, file.name))
            status = utils.safe_update_file(local_file.get_internal_path(),
                                            local_file.get_path(),
                                            file.get_content()
                                            )
            if status == 'Not changed':
                status = colors.info(status)
            elif status == 'Conflict':
                status = colors.error(status)
            else:
                status = colors.warning(status)
        else:
            status = colors.success('New')
            local_file = LocalFile()
            local_file.name = file.name.split('.')[0]
            local_file.dir = '' if flat else file.get_default_local_dir()
            local_file.type = file.type
            local_file.filename = file.name
            local_file.polygon_filename = file.name
            print('Downloading new file %s to %s' % (file.name, local_file.get_path()))
            content = file.get_content()
            utils.safe_rewrite_file(local_file.get_path(), content)
            utils.safe_rewrite_file(local_file.get_internal_path(), content)
            global_vars.problem.local_files.append(local_file)
        table.add_row([file.type, file.name, local_file.get_path(), status])
    print(table)
    save_session()


def process_add(args):
    if not load_session() or global_vars.problem.sessionId is None:
        fatal('No session known. Use relogin or init first.')
    if len(args) < 2:
        print_help()
    valid_types = ['solution', 'source', 'checker', 'validator']
    try:
        options, args = getopt.gnu_getopt(args, "t:")
    except getopt.GetoptError as err:
        fatal(err.msg)
    if args[0] not in valid_types:
        fatal('type should be in ' + str(valid_types))
    as_checker = False
    as_validator = False
    solution_type = None
    for opt in options:
        if opt[0] == '-t':
            solution_type = opt[1]
        else:
            pass # will be used with more options
    if args[0] == 'checker' or args[0] == 'validator':
        if len(args) != 2:
            fatal('can''t set several ' + args[0] + 's')
        if args[0] == 'checker':
            as_checker = True
        else:
            as_validator = True
        args[0] = 'source'
    if solution_type is not None:
        if args[0] != 'solution':
            fatal('solution type can be set only on solutions')
        valid_solution_types = ['MA', 'OK', 'RJ', 'TL', 'WA', 'PE', 'ML', 'RE']
        solution_type = solution_type.upper()
        solution_type_alias = {x: x for x in valid_solution_types}
        solution_type_alias.update({'MAIN': 'MA'})
        if solution_type not in solution_type_alias:
            fatal('solution type should be in ' + str(sorted(list(solution_type_alias.keys()))))
        solution_type = solution_type_alias[solution_type]
    table = PrettyTable(['File type', 'Polygon name', 'Local path', 'Status'])
    for filename in args[1:]:
        local = global_vars.problem.get_local_by_filename(os.path.basename(filename))
        if local is not None:
            print('file %s already added, use commit instead' % os.path.basename(filename))
            status = colors.warning('Already added')
        else:
            local = LocalFile(os.path.basename(filename),
                              os.path.dirname(filename),
                              str(os.path.basename(filename).split('.')[0]),
                              args[0]
                              )
            if local.upload():
                status = colors.success('Uploaded')
                global_vars.problem.local_files.append(local)
                if as_checker:
                    global_vars.problem.set_checker_validator(local.polygon_filename, 'checker')
                if as_validator:
                    global_vars.problem.set_checker_validator(local.polygon_filename, 'validator')
                if solution_type is not None:
                    global_vars.problem.change_solution_type(local.polygon_filename, solution_type.upper())
            else:
{'fraction_non_alphanumeric': 0.10356094958655641, 'fraction_numerical': 0.002800746865830888, 'mean_word_length': 2.3025765249944947, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6784476', 'n_tokens_mistral': 4361, 'n_tokens_neox': 4258, 'n_words': 1020}