setta 0.0.1.dev0__py3-none-any.whl → 0.0.2__py3-none-any.whl

setta/code_gen/python/make_parseable.py

@@ -0,0 +1,97 @@
+import ast
+import re
+
+
+def make_parseable(code: str, var_name: str) -> str:
+    """
+    Make minimal modifications to a potentially incomplete Python variable declaration
+    to make it parseable by ast.parse.
+
+    Args:
+        code (str): The potentially incomplete Python code
+
+    Returns:
+        str: Modified code that can be parsed by ast.parse
+    """
+    # No need to extract var_name as it's passed in
+    if is_parseable(code):
+        return code
+
+    code = code.strip()
+
+    # Strategy 1: Try removing the last character if it's a bracket or operator
+    if code[-1] in "([{=+-*/,":
+        modified = code[:-1]
+        if is_parseable(modified):
+            return modified
+
+    # Strategy 2: Try balancing parentheses/brackets/braces
+    opening = "([{"
+    closing = ")]}"
+    brackets_map = dict(zip(opening, closing))
+    stack = []
+
+    for char in code:
+        if char in opening:
+            stack.append(char)
+        elif char in closing:
+            if not stack or brackets_map[stack[-1]] != char:
+                # Mismatched bracket - remove this character
+                modified = code.replace(char, "", 1)
+                if is_parseable(modified):
+                    return modified
+            else:
+                stack.pop()
+
+    # Add missing closing brackets
+    while stack:
+        code += brackets_map[stack.pop()]
+        if is_parseable(code):
+            return code
+
+    # Strategy 3: Handle incomplete string literals
+    if re.search(r'[\'"]', code):
+        # Remove incomplete string if quote is the last character
+        if code[-1] in "\"'":
+            modified = code[:-1]
+            if is_parseable(modified):
+                return modified
+        # Try to close the string
+        for quote in ['"', "'"]:
+            modified = code + quote
+            if is_parseable(modified):
+                return modified
+
+    # Strategy 4: Remove trailing operators
+    operators = ["+", "-", "*", "/", "=", "&", "|", "^", "%"]
+    for op in operators:
+        if code.rstrip().endswith(op):
+            modified = code.rstrip()[:-1]
+            if is_parseable(modified):
+                return modified
+
+    # If nothing else works, try removing characters from the end until it parses
+    while len(code) > 0:
+        if is_parseable(code):
+            return code
+        code = code[:-1]
+
+    # this means it failed
+    return None
+
+
+def is_parseable(code: str) -> bool:
+    """
+    Check if the given code can be parsed by ast.parse
+
+    Args:
+        code (str): Python code to check
+
+    Returns:
+        bool: True if code can be parsed, False otherwise
+    """
+    try:
+        ast.parse(code)
+        return True
+    except SyntaxError:
+        return False

setta/code_gen/python/position_line_col.py

@@ -0,0 +1,33 @@
+def position_to_line_col(text, position):
+    if position < 0 or position > len(text):
+        return None
+
+    lines = text.split("\n")
+    current_pos = 0
+
+    for line_num, line in enumerate(lines):
+        line_length = len(line) + 1  # +1 for the newline character
+        if current_pos + line_length > position:
+            col = position - current_pos
+            return line_num, col
+        current_pos += line_length
+
+    # If we get here, the position is at the very end of the text
+    return len(lines) - 1, len(lines[-1])
+
+
+def line_col_to_position_batch(text, line_col_pairs):
+    lines = text.split("\n")
+    line_starts = [0]
+    for line in lines[:-1]:
+        line_starts.append(line_starts[-1] + len(line) + 1)
+
+    results = []
+    for line, col in line_col_pairs:
+        if line < 0 or line >= len(lines) or col < 0:
+            results.append(None)
+        else:
+            position = line_starts[line] + min(col, len(lines[line]))
+            results.append(position)
+
+    return results

setta/code_gen/python/validate_imports.py

@@ -0,0 +1,87 @@
+import asyncio
+import copy
+import logging
+
+from setta.code_gen.export_selected import get_specific_template_var
+from setta.code_gen.find_placeholders import tp
+from setta.database.utils import create_new_id
+from setta.utils.constants import C
+
+logger = logging.getLogger(__name__)
+
+
+async def maybe_validate_imports_and_update_code(
+    sectionId, code, ref_template_var_positions, lsp_writers
+):
+    templateVars = ref_template_var_positions["templateVars"]
+    generated_imports_info, idx = get_specific_template_var(
+        templateVars, tp(C.SETTA_GENERATED_PYTHON_IMPORTS), return_index=True
+    )
+    if not generated_imports_info:
+        return code, templateVars
+    old_value = generated_imports_info["value"]
+    new_value, did_change = await validate_imports(sectionId, old_value, lsp_writers)
+    if not did_change:
+        return code, templateVars
+
+    # this shouldn't happen
+    if len(new_value) != len(old_value):
+        logger.debug("new imports have different length from old imports")
+        return code, templateVars
+
+    templateVars = copy.deepcopy(ref_template_var_positions["templateVars"])
+    generated_imports_info = templateVars[idx]
+    generated_imports_info["value"] = new_value
+    startPos = generated_imports_info["startPos"]
+    endPos = generated_imports_info["startPos"] + len(new_value)
+    code = code[:startPos] + new_value + code[endPos:]
+    return code, templateVars
+
+
+async def validate_imports(sectionId, generated_imports, lsp_writers):
+    if not generated_imports:
+        return generated_imports, False
+
+    code_id = f"{sectionId}-validate-imports"
+    await lsp_writers["full"].send_document_open_request(
+        code_id=code_id, code=generated_imports, virtual=True
+    )
+
+    lines = generated_imports.split("\n")
+    message_id = create_new_id()
+
+    # First, send all requests and store the futures
+    request_futures = []
+    line_indices = []  # Keep track of which line each request is for
+    for idx, line in enumerate(lines):
+        if not line:
+            continue
+        request_id = f"{message_id}-{idx}"
+        # Send request without awaiting
+        lsp_writers["full"].create_asyncio_task(
+            lsp_writers["full"].send_definition_request(
+                idx, len(line) - 1, code_id, request_id, virtual=True
+            ),
+            request_id,
+            local=True,
+        )
+        # Store the future for waiting later
+        future = lsp_writers["full"].wait_for_response(request_id)
+        request_futures.append(future)
+        line_indices.append(idx)  # Store the original line index
+
+    # Then await all responses concurrently
+    responses = await asyncio.gather(*request_futures)
+
+    did_change = False
+    # Use line_indices to map responses back to correct lines
+    for response_idx, response in enumerate(responses):
+        original_idx = line_indices[response_idx]
+        if not response["data"]:
+            # replace with equal length comment so that we don't have
+            # to change any position information
+            lines[original_idx] = "#" * len(lines[original_idx])
+            did_change = True
+
+    generated_imports = "\n".join(lines)
+    return generated_imports, did_change

setta/code_gen/utils.py

@@ -0,0 +1,120 @@
+import copy
+
+
+def push_var_deep(objects, target_var_name):
+    def find_dependents(obj_list, var_name):
+        """Find all objects that directly or indirectly depend on the given object."""
+        dependents = set()
+        for obj in obj_list:
+            if var_name in obj["dependencies"]:
+                dependents.add(obj["var_name"])
+                dependents.update(find_dependents(obj_list, obj["var_name"]))
+        return dependents
+
+    dependents = find_dependents(objects, target_var_name)
+    everything_else = []
+    target_and_dependents = []
+    for obj in objects:
+        if obj["var_name"] != target_var_name and obj["var_name"] not in dependents:
+            everything_else.append(obj)
+        else:
+            target_and_dependents.append(obj)
+
+    return everything_else + target_and_dependents
+
+
+def process_refs(
+    code,
+    refs,
+    get_ref_var_name,
+    cursor_position=None,
+    templateVars=None,
+    get_template_var_replacement_value=None,
+):
+    positions = {
+        "refs": [],  # list of dicts, where each dict is {"keyword", "startPos", "value"}
+        "templateVars": [],  # list of dicts, where each tuple is {"keyword", "startPos", "value"}
+        "cursor": cursor_position,
+    }
+    if templateVars is None:
+        templateVars = []
+
+    refs = copy.deepcopy(refs)
+    templateVars = copy.deepcopy(templateVars)
+    refs_set = set(id(x) for x in refs)
+    to_process = sorted(refs + templateVars, key=lambda x: x["startPos"])
+    args = [
+        positions,
+        to_process,
+        refs_set,
+        get_ref_var_name,
+        get_template_var_replacement_value,
+    ]
+    code = process_refs_helper(code, *args, process_last_ones=False)
+    code = process_refs_helper(code, *args, process_last_ones=True)
+    return code, positions
+
+
+def process_refs_helper(
+    code,
+    positions,
+    to_process,
+    refs_set,
+    get_ref_var_name,
+    get_template_var_replacement_value,
+    process_last_ones,
+):
+    input_positions_len = {k: len(v) for k, v in positions.items() if k != "cursor"}
+
+    for i, x in enumerate(to_process):
+        if x.get("processLast", False) != process_last_ones:
+            continue
+
+        if id(x) in refs_set:
+            replacement_value = get_ref_var_name(
+                x["sectionId"], x["paramInfoId"], x["isArgsObj"]
+            )
+            type = "refs"
+        else:
+            replacement_value = get_template_var_replacement_value(code, x, positions)
+            type = "templateVars"
+
+        startPos = x["startPos"]
+        endPos = startPos + len(x["keyword"])
+        length_diff = len(replacement_value) - (endPos - startPos)
+
+        # Update positions for all future replacements
+        for future_x in to_process[i + 1 :]:
+            future_x["startPos"] += length_diff
+
+        # Update positions for all positions that already existed when this function was called.
+        update_existing_positions(positions, input_positions_len, startPos, length_diff)
+
+        if positions["cursor"] is not None:
+            positions["cursor"] = get_new_cursor_position(
+                positions["cursor"], length_diff, startPos, endPos
+            )
+
+        positions[type].append(
+            {"keyword": x["keyword"], "startPos": startPos, "value": replacement_value}
+        )
+        code = code[:startPos] + replacement_value + code[endPos:]
+
+    return code
+
+
+def update_existing_positions(positions, input_positions_len, startPos, length_diff):
+    for k, v in input_positions_len.items():
+        for i in range(v):
+            if startPos < positions[k][i]["startPos"]:
+                positions[k][i]["startPos"] += length_diff
+
+
+def get_new_cursor_position(cursor_position, length_diff, startPos, endPos):
+    # Update position if replacement occurs before the initial position
+    if endPos <= cursor_position:
+        # All changes before the position affect its offset
+        return cursor_position + length_diff
+    elif startPos < cursor_position < endPos:
+        return endPos + length_diff  # Reset position to after the new substring
+    return cursor_position

setta/code_gen/yaml/generate_yaml.py

@@ -0,0 +1,23 @@
+import yaml
+
+
+def generate_yaml(e):
+    transformed_dict = {}
+    for key, value in e.items():
+        if "value" in value and isinstance(value["value"], dict):
+            # Handle nested structure for keys like 'dataset' and 'algorithm'
+            callable_value = value["value"]["callable"]
+            params = value["value"].get("params", [])
+            param_dict = {}
+            for param in params:
+                param_value = e[param]["value"]
+                # Resolve references to other keys (like 'n_samples' for 'p0')
+                if param_value in e:
+                    param_value = e[param_value]["value"]
+                param_dict[e[param]["name"]] = param_value
+            transformed_dict[key] = {callable_value: param_dict}
+        elif not value["is_global"]:
+            # Directly use top-level keys and values
+            transformed_dict[key] = value["value"]
+
+    return yaml.safe_dump(transformed_dict)

setta/code_gen/yaml/section_dict.py

@@ -0,0 +1,93 @@
+import yaml
+
+
+def section_dict_to_yaml(exporter_objs, codeInfo, sectionVariant):
+    pinned = exporter_objs["pinned"].prettyOutput
+    unpinned = exporter_objs["unpinned"].prettyOutput
+    pinnedStr = raw_dump(pinned)
+    unpinnedStr = raw_dump(unpinned)
+
+    yamlValue = ""
+    if sectionVariant["selectedItem"]:
+        selectedItemName = codeInfo[sectionVariant["selectedItem"]]["name"]
+        yamlValue += f"{selectedItemName}\n---\n"
+
+    if len(pinned) > 0:
+        yamlValue += f"{pinnedStr}\n---\n"
+    elif yamlValue != "":
+        yamlValue += "\n---\n"
+    if len(unpinned) > 0:
+        yamlValue += unpinnedStr
+
+    return yamlValue
+
+
+def param_sweep_dict_to_yaml(exporter_obj):
+    return "\n---\n".join(raw_dump(x) for x in exporter_obj.prettyOutput)
+
+
+def run_group_dict_to_yaml(exporter_obj):
+    if len(exporter_obj.prettyOutput) == 0:
+        return "{}"
+    return yaml.dump(
+        exporter_obj.prettyOutput,
+        Dumper=RunGroupDumper,
+        explicit_end=False,
+        sort_keys=False,
+    )
+
+
+# Register representer for dict that skips empty ones
+def represent_dict_skip_empty(dumper, data):
+    if not data:  # if dictionary is empty
+        return dumper.represent_scalar("tag:yaml.org,2002:null", "")
+    return dumper.represent_dict(data)
+
+
+class RunGroupDumper(yaml.Dumper):
+    pass
+
+
+RunGroupDumper.add_representer(dict, represent_dict_skip_empty)
+
+
+def raw_dump(input_data, indent=0):
+    output = ""
+
+    # Handle list of dictionaries or other lists
+    if isinstance(input_data, list):
+        for item in input_data:
+            output += " " * indent + "-"
+            if isinstance(item, (dict, list)):
+                output += " " + raw_dump(item, indent + 2)[indent + 2 :]
+            else:
+                output += " " + str(item) + "\n"
+    # Handle single dictionary
+    elif isinstance(input_data, dict):
+        for k, v in input_data.items():
+            # Add indentation
+            output += " " * indent
+
+            # If value is a dictionary, recursively process it
+            if isinstance(v, dict):
+                output += f"{k}:\n"
+                output += raw_dump(v, indent + 2)
+            else:
+                # Check if value is a list
+                if isinstance(v, list):
+                    output += f"{k}:\n"
+                    output += raw_dump(v, indent + 2)
+                # Check if value contains newlines
+                else:
+                    str_value = str(v)
+                    if "\n" in str_value:
+                        # Use literal block style
+                        output += f"{k}: |\n"
+                        # Indent each line of the value
+                        for line in str_value.splitlines():
+                            output += " " * (indent + 2) + line + "\n"
+                    else:
+                        # Single line value
+                        output += f"{k}: {str_value}\n"
+
+    return output

setta/database/backup.py

@@ -0,0 +1,80 @@
+import logging
+import os
+import shutil
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import List, Optional
+
+from setta.utils.constants import DB_BACKUP_FOLDER, USER_SETTINGS
+
+logger = logging.getLogger(__name__)
+
+
+def ensure_backup_dir_exists():
+    """Create the backup directory if it doesn't exist."""
+    Path(DB_BACKUP_FOLDER).mkdir(parents=True, exist_ok=True)
+
+
+def get_backup_files() -> List[Path]:
+    """Get a list of existing backup files, sorted by filename (which contains the timestamp)."""
+    backup_dir = Path(DB_BACKUP_FOLDER)
+    return sorted(
+        [f for f in backup_dir.glob("backup_*.db")], key=lambda x: x.name, reverse=True
+    )
+
+
+def get_last_backup_time() -> Optional[float]:
+    """
+    Get the timestamp of the last backup from the most recent backup filename.
+    Returns None if no backups exist.
+    """
+    backup_files = get_backup_files()
+    if not backup_files:
+        return None
+
+    latest_backup = backup_files[0]
+    timestamp_str = "_".join(latest_backup.stem.split("_")[-2:])
+
+    # Convert the timestamp string to a datetime object
+    backup_datetime = datetime.strptime(timestamp_str, "%Y%m%d_%H%M%S")
+
+    # Convert the datetime to a Unix timestamp (seconds since epoch)
+    return backup_datetime.timestamp()
+
+
+def cleanup_old_backups():
+    backup_files = get_backup_files()
+    for old_backup in backup_files[USER_SETTINGS["backend"]["autobackupKeepLastN"] :]:
+        old_backup.unlink()
+
+
+def create_backup(db_path):
+    """Create a new backup of the database."""
+    ensure_backup_dir_exists()
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    backup_filename = f"backup_{timestamp}.db"
+    backup_path = os.path.join(DB_BACKUP_FOLDER, backup_filename)
+    logger.debug(f"creating backup: {backup_path}")
+
+    shutil.copy2(db_path, backup_path)
+    cleanup_old_backups()
+
+
+def should_create_backup() -> bool:
+    """Check if it's time to create a new backup."""
+    if not USER_SETTINGS["backend"]["autobackup"]:
+        return False
+    last_backup_time = get_last_backup_time()
+    logger.debug(f"last_backup_time: {last_backup_time}")
+    if last_backup_time is None:
+        return True
+
+    current_time = time.time()
+    time_since_last_backup = current_time - last_backup_time
+    return time_since_last_backup >= USER_SETTINGS["backend"]["autobackupFreq"]
+
+
+def maybe_create_backup(db_path):
+    if should_create_backup():
+        create_backup(db_path)

setta/database/db/artifacts/load.py

@@ -0,0 +1,93 @@
+import json
+
+from setta.cli.logger import get_wrapper_for_loading
+from setta.utils.constants import C
+from setta.utils.utils import try_json
+
+
+def load_artifacts(db, ids):
+    placeholders = ", ".join(["?"] * len(ids))
+    query = f"""
+        SELECT id, name, path, value, type
+        FROM Artifact
+        WHERE id IN ({placeholders})
+    """
+
+    db.execute(query, ids)
+
+    output = {}
+    for row in db.fetchall():
+        output[row["id"]] = {
+            "name": row["name"],
+            "path": row["path"],
+            "value": try_json(row["value"]),
+            "type": row["type"],
+        }
+
+    return output
+
+
+def load_artifact_groups(db, ids):
+    placeholders = ", ".join(["?"] * len(ids))
+    query = f"""
+        SELECT ArtifactGroupId.id, ArtifactGroupId.name, ArtifactGroupId.data, ArtifactGroup.artifactId, ArtifactGroup.data as artifactTransformData, ArtifactGroup."order"
+        FROM ArtifactGroupId
+        LEFT JOIN ArtifactGroup
+        ON ArtifactGroupId.id = ArtifactGroup.idid
+        WHERE ArtifactGroupId.id in ({placeholders})
+        ORDER BY ArtifactGroupId.id, ArtifactGroup."order"
+    """
+    db.execute(query, ids)
+    output = {}
+    for row in db.fetchall():
+        if row["id"] not in output:
+            output[row["id"]] = {
+                "name": row["name"],
+                "artifactTransforms": [],
+                **json.loads(row["data"]),
+            }
+        if row["artifactId"]:
+            output[row["id"]]["artifactTransforms"].append(
+                {
+                    "artifactId": row["artifactId"],
+                    **json.loads(row["artifactTransformData"]),
+                }
+            )
+    return output
+
+
+def load_available_artifacts(db, sectionType):
+    allowed_types = C.ALLOWED_ARTIFACT_TYPES[sectionType]
+
+    placeholders = ", ".join(["?"] * len(allowed_types))
+    query = f"""
+        SELECT id, name, path, type
+        FROM Artifact
+        WHERE type IN ({placeholders})
+    """
+
+    db.execute(query, allowed_types)
+
+    output = {}
+    for row in db.fetchall():
+        output[row["id"]] = {
+            "name": row["name"],
+            "path": row["path"],
+            "type": row["type"],
+        }
+
+    return output
+
+
+def load_artifact_metadata_and_maybe_value_from_disk(db, artifactIds):
+    artifacts = load_artifacts(db, artifactIds)
+    for a in artifacts.values():
+        if a["path"]:
+            a.update(load_artifact_from_disk(a["path"], a["type"]))
+    return artifacts
+
+
+def load_artifact_from_disk(filepath, type):
+    wrapper = get_wrapper_for_loading(type)
+    wrapper.load(filepath)
+    return wrapper.serialize()