iris-hep / func_adl_xAOD

Client interface to send a hierarchical SQL-like query to an xAOD backend
0 stars 5 forks source link

SHould be able to deal with the `is` operator #224

Open gordonwatts opened 7 months ago

gordonwatts commented 7 months ago

In the following, at the very end, we should be able to do is True - right now we have to do == True

@pytest.mark.skipif(not run_long_running_tests, reason="Skipping long running test")
def test_class_enum_use_inline():
    """Use a enum as part of the actual query"""

    # Define the calo clusters return
    from func_adl_xAOD.atlas.xaod.event_collections import (
        atlas_xaod_collections,
    )

    atlas_xaod_collections.append(
        EventCollectionSpecification(
            "atlas",
            "CaloClusters",
            ["xAODCaloEvent/CaloClusterContainer.h"],
            atlas_xaod_event_collection_collection(
                "xAOD::CaloClusterContainer", "xAOD::CaloCluster"
            ),
            ["xAODCaloEvent"],
        ),
    )

    # Define the num class for future reference
    class xAOD:
        class CaloCluster_v1:
            class ClusterSize(Enum):
                SW_55ele = 1
                SW_35ele = 2
                SW_37ele = 3
                SW_55gam = 4
                SW_35gam = 5
                SW_37gam = 6
                SW_55Econv = 7
                SW_35Econv = 8
                SW_37Econv = 9
                SW_softe = 10
                Topo_420 = 11
                Topo_633 = 12
                SW_7_11 = 13
                SuperCluster = 14
                Tower_01_01 = 15
                Tower_005_005 = 16
                Tower_fixed_area = 17
                CSize_Unknown = 99

    # Declare the enum class
    ctyp.define_enum(
        "xAOD.CaloCluster_v1",
        "ClusterSize",
        [e.name for e in xAOD.CaloCluster_v1.ClusterSize],
    )

    # Define the `getConstituentsSignalState` method
    ctyp.add_method_type_info(
        "xAOD::CaloCluster_v1",
        "clusterSize",
        ctyp.terminal("xAOD::CaloCluster_v1::ClusterSize", tree_type="int"),
    )

    # fmt: off
    training_df = as_pandas(
        f_exot_15.SelectMany(lambda e: e.CaloClusters("egammaClusters"))
        .Select(lambda c: c.clusterSize() == xAOD.CaloCluster_v1.ClusterSize.SuperCluster)
    )
    # fmt: on

    assert len(training_df) > 0
    print(training_df.col1)
    assert all(training_df.col1 is True)