# Copyright (c) 2024-2026 Karana Dynamics Pty Ltd. All rights reserved.
#
# NOTICE TO USER:
#
# This source code and/or documentation (the "Licensed Materials") is
# the confidential and proprietary information of Karana Dynamics Inc.
# Use of these Licensed Materials is governed by the terms and conditions
# of a separate software license agreement between Karana Dynamics and the
# Licensee ("License Agreement"). Unless expressly permitted under that
# agreement, any reproduction, modification, distribution, or disclosure
# of the Licensed Materials, in whole or in part, to any third party
# without the prior written consent of Karana Dynamics is strictly prohibited.
#
# THE LICENSED MATERIALS ARE PROVIDED "AS IS" WITHOUT WARRANTY OF ANY KIND.
# KARANA DYNAMICS DISCLAIMS ALL WARRANTIES, EXPRESS OR IMPLIED, INCLUDING
# BUT NOT LIMITED TO WARRANTIES OF MERCHANTABILITY, NON-INFRINGEMENT, AND
# FITNESS FOR A PARTICULAR PURPOSE.
#
# IN NO EVENT SHALL KARANA DYNAMICS BE LIABLE FOR ANY DAMAGES WHATSOEVER,
# INCLUDING BUT NOT LIMITED TO LOSS OF PROFITS, DATA, OR USE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGES, WHETHER IN CONTRACT, TORT,
# OR OTHERWISE ARISING OUT OF OR IN CONNECTION WITH THE LICENSED MATERIALS.
#
# U.S. Government End Users: The Licensed Materials are a "commercial item"
# as defined at 48 C.F.R. 2.101, and are provided to the U.S. Government
# only as a commercial end item under the terms of this license.
#
# Any use of the Licensed Materials in individual or commercial software must
# include, in the user documentation and internal source code comments,
# this Notice, Disclaimer, and U.S. Government Use Provision.
import os
import re
from contextlib import contextmanager
from typing import runtime_checkable, Protocol, TYPE_CHECKING, overload, cast
from pint import Quantity
import numpy as np
import Karana.Core as kc
from pathlib import Path
if TYPE_CHECKING:
from Karana.KUtils.DataStruct import DataStruct, IdMixin
else:
from .DataStruct import DataStruct, IdMixin
@runtime_checkable
class _Named(Protocol):
name: str # Any class matching this protocol must have a 'name' attribute that is a string
[docs]
class VarPythonScriptContext:
def __init__(self):
self._names_dict: dict[object | int, str] = {}
# Add reserved names
self._names: set[str] = {"o_sh", "p_sh", "sh", "vroot"}
# Common regexs used
self._name_match_re = re.compile("[a-zA-Z_][a-zA-Z0-9_]+")
self._rename_re = re.compile(r"(?<=[a-z])(?=[A-Z])|[^a-zA-Z0-9]")
@overload
def getVarNameIfExists(self, ds: DataStruct) -> str | None: ...
@overload
def getVarNameIfExists(self, obj: object) -> str | None: ...
[docs]
def getVarNameIfExists(self, *args, **_):
if isinstance(ds := args[0], DataStruct):
# If this is already in the cache, then return it
if body_name := self._names_dict.get(id(ds)):
# Search by id
return body_name
elif isinstance(ds, IdMixin) and len(ds._objects_from_id) == 1:
# Search by Base object if it exists
obj = cast(kc.Base | None, ds._objects_from_id[0]())
if obj is not None and (obj_name := self._names_dict.get(obj, None)):
# If found by object, add by ID as well
self._names_dict[id(ds)] = obj_name
# Return the object
return obj_name
else:
return None
else:
obj = cast(object, args[0])
# If this is already in the names cache, then return it
return self._names_dict.get(obj, None)
@overload
def getVarName(self, ds: DataStruct, base_name: str = "var") -> str: ...
@overload
def getVarName(self, obj: object, base_name: str = "var") -> str: ...
[docs]
def getVarName(self, *args, **kwargs):
# Return the name if it already exists
if name := self.getVarNameIfExists(*args, **kwargs):
return name
# Otherwise, create a new name
if isinstance(ds := args[0], DataStruct):
base_name = kwargs.get("base_name", "var")
if isinstance(ds, _Named) and re.match(self._name_match_re, ds.name):
# If the name of the body can become a variable, then use that
dark = re.sub(self._rename_re, "_", ds.name).lower()
else:
# Otherwise, just use the base_name
dark = base_name
# Add on a number to make this unique if necessary
body_name = dark
k = 0
while body_name in self._names:
k += 1
body_name = dark + f"_{k}"
# Add the new name to the cache
self._names.add(body_name)
self._names_dict[id(ds)] = body_name
if isinstance(ds, IdMixin) and len(ds._objects_from_id) == 1:
body = ds._objects_from_id[0]()
if body is not None:
self._names_dict[body] = body_name
return body_name
else:
base_name = kwargs.get("base_name", "var")
if isinstance(base := args[0], kc.Base):
base = cast(kc.Base, args[0])
if re.match(self._name_match_re, base.name()):
# If the name of the body can become a variable, then use that
dark = re.sub(self._rename_re, "_", base.name()).lower()
else:
# Otherwise, just use the base_name
dark = base_name
else:
dark = base_name
# Add on a number to make this unique if necessary
obj_name = dark
k = 0
while obj_name in self._names:
k += 1
obj_name = dark + f"_{k}"
# Add the new name to the cache
self._names.add(obj_name)
self._names_dict[base] = obj_name
return obj_name
[docs]
def valuePrinter(q) -> str:
"""Print values for Python scripts.
Parameters
----------
q : Quantity | NDArray
The value to print
Returns
-------
str
The string representation of the value.
"""
if isinstance(q, Quantity):
if isinstance(q.m, np.ndarray):
return "np." + repr(q.m) + " * " + re.sub(r"([a-z]+)", r"ureg.\1", str(q.u))
elif isinstance(q.m, float):
return repr(q.m) + " * " + re.sub(r"([a-z]+)", r"ureg.\1", str(q.u))
else:
raise ValueError(f"Not sure how to print out quantity with magnitude type {type(q.m)}.")
elif isinstance(q, np.ndarray):
return "np." + repr(q)
elif isinstance(q, float):
return repr(q)
else:
raise ValueError("Type unknown.")
[docs]
@contextmanager
def fullPrintOpts():
"""Context to print numpy values with full precision."""
opts = np.get_printoptions()
try:
np.set_printoptions(floatmode="unique")
yield
finally:
np.set_printoptions(**opts)
def writePythonScript(code: str, file: Path):
"""Write the Python script to the provided Path.
This adds comment seperators with the appropriate number of -'s.
It also runs ruff on the output.
"""
comment_sep = re.compile(r"([ \t]+)_KARANA_COMMENT_SEPARATOR")
def replaceCommentSep(match: re.Match) -> str:
"""Replace the incoming match with a comment separator.
Parameters
----------
match : Math
The regex match to fix.
Returns
-------
str
The comment separator.
"""
indent = match.group(1)
prefix = "# "
dashes = "-" * (100 - len(indent) - len(prefix))
return f"{indent}{prefix}{dashes}"
# Replace comment separators
code = re.sub(comment_sep, replaceCommentSep, code)
# Write the file
with open(file, "w") as f:
f.write(code)
# Format the file with ruff
wd = file.parent
os.system(f"cd {wd} && ruff format --config 'line-length=100' {file}")
