From 87e7156d8f1d3217d321b336fe95f627014de173 Mon Sep 17 00:00:00 2001 From: fscarlier Date: Wed, 26 Jul 2023 16:27:09 +0200 Subject: [PATCH] added boolean %b to headers for MADNG output --- tests/inputs/bool_header.tfs | 25 +++++++++++++++++++++++++ tests/test_reader.py | 16 ++++++++++++++++ tfs/constants.py | 1 + tfs/reader.py | 12 +++++++++++- tfs/writer.py | 6 +++++- 5 files changed, 58 insertions(+), 2 deletions(-) create mode 100644 tests/inputs/bool_header.tfs diff --git a/tests/inputs/bool_header.tfs b/tests/inputs/bool_header.tfs new file mode 100644 index 00000000..336e328b --- /dev/null +++ b/tests/inputs/bool_header.tfs @@ -0,0 +1,25 @@ +@ TITLE %s "Title of your tfs file" +@ DPP %le 1.0 +@ Q1 %le 0.269974877952 +@ Q1RMS %le 1.75642567736e-07 +@ NATQ1 %le 0.280041400187 +@ NATQ1RMS %le 0.00102479265802 +@ BPMCOUNT %d 9 +@ BOOLTRUE1 %b true +@ BOOLTRUE2 %b True +@ BOOLTRUE3 %b 1 +@ BOOLFALSE1 %b false +@ BOOLFALSE2 %b False +@ BOOLFALSE3 %b 0 +# Some comment line here +* NAME S NUMBER CO CORMS BPM_RES +$ %s %le %d %le %le %le + BPMYB.5L2.B1 28.288 1 -0.280727353099 0.00404721900879 0.121264541395 + BPMYB.4L2.B1 48.858 2 0.601472827003 0.00301396244054 0.129738519811 + BPMWI.4L2.B1 73.3255 3 -0.610294990396 0.0039123010318 0.0952864848273 + BPMSX.4L2.B1 123.4825 3472136972 0.778206651453 0.00542543379504 0.0578581425476 + "BPMS.2L2.B1" 161.394 59055944 0.585105573645 0.00291016910226 0.1223625619 + "BPMSW.1L2.B1" 171.328 09202215 2.50235465023 0.00275350035218 0.148603785488 + BPMSW.1R2.B1 214.518 3117 1.81036167087 0.00282138482457 0.164954082556 + BPMS.2R2.B1 224.452 18943819309 0.0791371365672 0.00474290041487 0.122265653712 + BPMSX.4R2.B1 262.3635 105 -0.00665768479832 0.00350302654669 0.187320306406 diff --git a/tests/test_reader.py b/tests/test_reader.py index 224d697e..c088f9c0 100644 --- a/tests/test_reader.py +++ b/tests/test_reader.py @@ -137,6 +137,17 @@ def test_read_file_without_header(self, _tfs_file_without_header, _tfs_file_path df_for_compare = read_tfs(_tfs_file_pathlib) assert_frame_equal(df, df_for_compare) + def test_read_file_with_bool_in_header(self, _tfs_file_with_bool_in_header, _tfs_file_pathlib): + df = read_tfs(_tfs_file_with_bool_in_header) + assert df.headers['BOOLTRUE1'] == True + assert df.headers['BOOLTRUE2'] == True + assert df.headers['BOOLTRUE3'] == True + assert df.headers['BOOLFALSE1'] == False + assert df.headers['BOOLFALSE2'] == False + assert df.headers['BOOLFALSE3'] == False + df_for_compare = read_tfs(_tfs_file_pathlib) + assert_frame_equal(df, df_for_compare) + class TestFailures: def test_absent_attributes_and_keys(self, _tfs_file_str: str): @@ -235,6 +246,11 @@ def _tfs_file_without_header() -> pathlib.Path: return INPUTS_DIR / "no_header.tfs" +@pytest.fixture() +def _tfs_file_with_bool_in_header() -> pathlib.Path: + return INPUTS_DIR / "bool_header.tfs" + + @pytest.fixture() def _tfs_file_without_header_but_empty_line() -> pathlib.Path: return INPUTS_DIR / "no_header_just_an_empty_line.tfs" diff --git a/tfs/constants.py b/tfs/constants.py index 5e08769f..1b341b57 100644 --- a/tfs/constants.py +++ b/tfs/constants.py @@ -19,6 +19,7 @@ ID_TO_TYPE: Dict[str, type] = { # used when reading files "%s": str, + "%b": bool, "%bpm_s": str, "%le": np.float64, "%f": np.float64, diff --git a/tfs/reader.py b/tfs/reader.py index 27f2b55e..ebb4c6a7 100644 --- a/tfs/reader.py +++ b/tfs/reader.py @@ -283,7 +283,10 @@ def _parse_header(str_list: List[str]) -> tuple: name = " ".join(str_list[0:type_index]) value_string = " ".join(str_list[(type_index + 1) :]) - return name, _id_to_type(str_list[type_index])(value_string.strip('"')) + if _id_to_type(str_list[type_index]) == bool: + return name, _string_to_bool(value_string.strip('"')) + else: + return name, _id_to_type(str_list[type_index])(value_string.strip('"')) def _find_and_set_index(data_frame: TfsDataFrame) -> TfsDataFrame: @@ -311,6 +314,13 @@ def _compute_types(str_list: List[str]) -> List[type]: return [_id_to_type(string) for string in str_list] +def _string_to_bool(string: str) -> bool: + if string.lower() in ['true', '1']: + return True + elif string.lower() in ['false', '0']: + return False + + def _id_to_type(type_str: str) -> type: try: return ID_TO_TYPE[type_str] diff --git a/tfs/writer.py b/tfs/writer.py index af51cfd5..f89bcb64 100644 --- a/tfs/writer.py +++ b/tfs/writer.py @@ -194,6 +194,8 @@ def _get_header_line(name: str, value, width: int) -> str: type_str = _value_to_type_string(value) if type_str == "%s": value = f'"{value}"' + if type_str == "%b": + value = str(bool(value)) return f"@ {name:<{width}} {type_str} {value:>{width}}" @@ -256,10 +258,12 @@ def _dtype_to_id_string(type_: type) -> str: Returns: The ID string. """ - if pdtypes.is_integer_dtype(type_) or pdtypes.is_bool_dtype(type_): + if pdtypes.is_integer_dtype(type_): return "%d" elif pdtypes.is_float_dtype(type_): return "%le" + elif pdtypes.is_bool_dtype(type_): + return "%b" elif pdtypes.is_string_dtype(type_): return "%s" raise TypeError(