Skip to content

Commit

Permalink
tidy, types
Browse files Browse the repository at this point in the history
  • Loading branch information
nickzoic committed Aug 26, 2024
1 parent 9a2f931 commit 69b9cc2
Show file tree
Hide file tree
Showing 4 changed files with 45 additions and 13 deletions.
15 changes: 12 additions & 3 deletions countess/core/parameters.py
Original file line number Diff line number Diff line change
Expand Up @@ -455,7 +455,7 @@ def get_column(self, df):


class ColumnOrStringParam(ColumnChoiceParam):
DEFAULT_VALUE = ""
DEFAULT_VALUE : Any = ""
PREFIX = "— "

def set_column_choices(self, choices):
Expand All @@ -475,7 +475,7 @@ def get_value(self, data: dict):
def get_column_or_value(self, df: pd.DataFrame, numeric: bool):
if self.value.startswith(self.PREFIX):
col = df[self.value[len(self.PREFIX) :]]
return col.astype('f' if numeric else 'string')
return col.astype("f" if numeric else "string")
else:
return float(self.value) if numeric else str(self.value)

Expand All @@ -493,7 +493,16 @@ def set_choices(self, choices: Iterable[str]):


class ColumnOrIntegerParam(ColumnOrStringParam):
DEFAULT_VALUE = 0
DEFAULT_VALUE : int = 0

def __init__(
self,
label: str,
value: Optional[int] = 0,
choices: Optional[Iterable[str]] = None,
):
super().__init__(label, choices=choices)
self.value = value

def clean_value(self, value):
if type(value) is str and value.startswith(self.PREFIX):
Expand Down
11 changes: 10 additions & 1 deletion countess/plugins/variant.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,14 @@

from countess import VERSION
from countess.core.logger import Logger
from countess.core.parameters import BooleanParam, ColumnChoiceParam, ColumnOrIntegerParam, ColumnOrStringParam, IntegerParam, StringParam
from countess.core.parameters import (
BooleanParam,
ColumnChoiceParam,
ColumnOrIntegerParam,
ColumnOrStringParam,
IntegerParam,
StringParam,
)
from countess.core.plugins import PandasTransformDictToDictPlugin
from countess.utils.variant import find_variant_string

Expand All @@ -31,6 +38,8 @@ class VariantPlugin(PandasTransformDictToDictPlugin):

def process_dict(self, data, logger: Logger) -> dict:
assert isinstance(self.parameters["reference"], ColumnOrStringParam)
assert isinstance(self.parameters["offset"], ColumnOrIntegerParam)

sequence = data[self.parameters["column"].value]
if not sequence:
return {}
Expand Down
26 changes: 18 additions & 8 deletions tests/plugins/test_python.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import pandas as pd
import numpy as np
import pandas as pd
import pytest

from countess.core.logger import ConsoleLogger
Expand All @@ -8,16 +8,23 @@
logger = ConsoleLogger()


dfi = pd.DataFrame([[1,1,2,7,11],[2,2,3,8,1],[3,3,4,9,3],[4,4,5,10,4],[5,5,6,12,7]], columns=['a','b','c','d','e'])
dfi = pd.DataFrame(
[[1, 1, 2, 7, 11], [2, 2, 3, 8, 1], [3, 3, 4, 9, 3], [4, 4, 5, 10, 4], [5, 5, 6, 12, 7]],
columns=["a", "b", "c", "d", "e"],
)


def test_python_builtins():
plugin = PythonPlugin()
plugin.set_parameter("code", """
plugin.set_parameter(
"code",
"""
x = mean(a,b,c,d,e)
y = std(a,b,c,d,e)
z = sqrt(pow(a,2) + pow(b,2))
v = var(a,b,c,d,e)
""")
""",
)

plugin.prepare(["test"], None)
dfo = plugin.process_dataframe(dfi, logger)
Expand All @@ -31,11 +38,14 @@ def test_python_builtins():

def test_python_dropna():
plugin = PythonPlugin()
plugin.set_parameter("code", """
plugin.set_parameter(
"code",
"""
a = None
n = None
if d >= 10: d = None
""")
""",
)
plugin.set_parameter("dropna", True)

plugin.prepare(["test"], None)
Expand All @@ -45,5 +55,5 @@ def test_python_dropna():
assert "n" not in dfo.columns
assert "d" in dfo.columns

assert any(np.isnan(dfo['d']))
assert not any(np.isnan(dfo['b']))
assert any(np.isnan(dfo["d"]))
assert not any(np.isnan(dfo["b"]))
6 changes: 5 additions & 1 deletion tests/plugins/test_variant.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,11 @@ def test_variant_ref_column():

def test_variant_ref_offset():
input_df = pd.DataFrame(
[{"seq": "TGAAGTAGAGG", "offs": "0"}, {"seq": "AGAAGTTGTGG", "offs": "10"}, {"seq": "ATAAGAAGAGG", "offs": "100"}]
[
{"seq": "TGAAGTAGAGG", "offs": "0"},
{"seq": "AGAAGTTGTGG", "offs": "10"},
{"seq": "ATAAGAAGAGG", "offs": "100"},
]
)

plugin = VariantPlugin()
Expand Down

0 comments on commit 69b9cc2

Please sign in to comment.