Skip to content

Commit

Permalink
Merge branch 'main' into nt/gh-207
Browse files Browse the repository at this point in the history
  • Loading branch information
ntessore committed Nov 25, 2024
2 parents 7769902 + 779246d commit 89a5b4f
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 27 deletions.
39 changes: 25 additions & 14 deletions heracles/io.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
import logging
import os
import re
from collections.abc import MutableMapping
from collections.abc import MutableMapping, Sequence
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Union
Expand Down Expand Up @@ -76,25 +76,36 @@ def _string_from_key(key: _DictKey) -> str:
"""
Return string representation for a given key.
"""
if isinstance(key, tuple):
names = list(map(_string_from_key, key))
c = ";" if any("," in name for name in names) else ","
return c.join(names)
return re.sub(r"\W+", "_", str(key))
# recursive expansion for sequences
if isinstance(key, Sequence) and not isinstance(key, str):
return "-".join(map(_string_from_key, key))

# get string representation of key
s = str(key)

# escape literal "\"
s = s.replace("\\", "\\\\")

# escape literal "-"
s = s.replace("-", "\\-")

# substitute non-FITS characters by tilde
s = re.sub(r"[^ -~]+", "~", s, flags=re.ASCII)

return s


def _key_from_string(s: str) -> _DictKey:
"""
Return key for a given string representation.
"""
keys = s.split(";")
if len(keys) > 1:
return tuple(map(_key_from_string, keys))
keys = keys[0].split(",")
if len(keys) > 1:
return tuple(map(_key_from_string, keys))
key = keys[0]
return int(key) if key.isdigit() else key
parts = re.split(r"(?<!\\)-", s.replace("\\\\", "\0"))
if len(parts) > 1:
return tuple(map(_key_from_string, parts))
key = parts[0]
key = key.replace("\\-", "-")
key = key.replace("\0", "\\")
return int(key) if key.removeprefix("-").isdigit() else key


def _get_next_extname(fits, prefix):
Expand Down
35 changes: 22 additions & 13 deletions tests/test_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -189,25 +189,32 @@ def test_string_from_key():
assert _string_from_key(1) == "1"
assert _string_from_key(("a",)) == "a"
assert _string_from_key((1,)) == "1"
assert _string_from_key(("a", 1)) == "a,1"
assert _string_from_key(("a", "b", 1, 2)) == "a,b,1,2"
assert _string_from_key((("a", 1), "b")) == "a,1;b"
assert _string_from_key((("a", 1), ("b", 2))) == "a,1;b,2"
assert _string_from_key(("a", 1)) == "a-1"
assert _string_from_key(("a", "b", 1, 2)) == "a-b-1-2"

# flatten nested sequences
assert _string_from_key([("a", 1), "b"]) == "a-1-b"
assert _string_from_key([("a", 1), ("b", (2,))]) == "a-1-b-2"

# test special chars
assert _string_from_key("a,b,c") == "a_b_c"
assert _string_from_key("!@#$%^&*()[]{};,.") == "_"
assert _string_from_key("a-b-c") == r"a\-b\-c"
assert _string_from_key(("a\\", 1)) == r"a\\-1"
assert _string_from_key(("a\\-", 1)) == r"a\\\--1"
assert _string_from_key(("a\\", -1)) == r"a\\-\-1"
assert _string_from_key("a€£") == "a~"


def test_key_from_string():
from heracles.io import _key_from_string

assert _key_from_string("a") == "a"
assert _key_from_string("1") == 1
assert _key_from_string("a,1") == ("a", 1)
assert _key_from_string("a,b,1,2") == ("a", "b", 1, 2)
assert _key_from_string("a,1;b") == (("a", 1), "b")
assert _key_from_string("a,1;b,2") == (("a", 1), ("b", 2))
assert _key_from_string("a-1") == ("a", 1)
assert _key_from_string("a-b-1-2") == ("a", "b", 1, 2)
assert _key_from_string(r"a\-b\-c") == "a-b-c"
assert _key_from_string(r"a\\-1") == ("a\\", 1)
assert _key_from_string(r"a\\\-1") == "a\\-1"
assert _key_from_string(r"a\\-\-1") == ("a\\", -1)


def test_write_read_maps(rng, tmp_path):
Expand Down Expand Up @@ -325,9 +332,11 @@ def test_write_read_cov(mock_cls, tmp_path):
workdir = str(tmp_path)

cov = {}
for k1, k2 in combinations_with_replacement(mock_cls, 2):
cl1, cl2 = mock_cls[k1], mock_cls[k2]
cov[k1, k2] = np.outer(cl1, cl2)
for (a1, b1, i1, j1), (a2, b2, i2, j2) in combinations_with_replacement(
mock_cls, 2
):
cl1, cl2 = mock_cls[a1, b1, i1, j1], mock_cls[a2, b2, i2, j2]
cov[a1, b1, a2, b2, i1, j1, i2, j2] = np.outer(cl1, cl2)

filename = "cov.fits"

Expand Down

0 comments on commit 89a5b4f

Please sign in to comment.