Any way to write custom checks that can be applied to any column?

[EvavW]

I'd like to write a special version of the isin check that can take a dictionary of the possible int values for a field, along with the string values for each enum. I would want to be able to apply this check to multiple fields in a schema, without defining it separately for each. May look roughly like:

from typing import Optional

import pandas as pd
import pandera as pa

from pandera.typing import Series
from pandera import SchemaModel, Field

class SCREEN(SchemaModel):
    class Config():
        code_book = {
            "sex": {0: "Male", 1: "Fema"},
            "chldbear": {0: "Unable", 1: "Able"},
        }

    @pa.check()
    def int_isin(cls, col: Series[int], code_book: dict) -> bool:
        return set(col.unique()) == Set(code_book.keys())

    sex: Series[int] = Field(int_isin=Config.code_book["sex"])
    chldbear: Optional[Series[int]] = Field(int_isin=Config.code_book["chldbear"])
    age: Option[Series[int]]

In this example I could define int_isin generally and then use it for any Field. Is this possible currently? If you know of a better way to reference Enums with both int and string values that would eliminate the need for this I would be happy to hear about that as well!

[jeffzi]

Hi @EvavW

In this example I could define int_isin generally and then use it for any Field. Is this possible currently?

You can apply a check to all columns with the regex argument: @pa.check(".*", regex=True)

Another trick is to access the series's name to look up the allowed values.

Here is how I would go about it:

from typing import Optional
import pandas as pd
import pandera as pa
from pandera.typing import Series


class Screen(pa.SchemaModel):
  # You don't need a Config, private attributes cannot be fields.
   _code_book = {
      "sex": {0: "Male", 1: "Fema"},
      "chldbear": {0: "Unable", 1: "Able"},
   }

   sex: Series[int]
   chldbear: Optional[Series[int]]
   age: Optional[Series[int]]

   @pa.check(".*", regex=True)
   @classmethod
   def check_codes(cls, series: Series) -> bool:
      col_name = series.name
      # look up allowed values if any
      codes = cls._code_book.get(col_name, {}).keys()
      if codes:
         return series.isin(codes)
      return True # nothing to check

df = pd.DataFrame({"sex": [1], "chldbear": [9], "age": [18]})
Screen.validate(df)

Note that in your example, Field(int_isin=Config.code_book["sex"]) will not work.
You need to register your check first (outside of the class). See doc here.

If you know of a better way to reference Enums with both int and string values

I'm not sure to understand that part. Did I answer all your questions?

[EvavW]

Thanks for the response! This is a nice solution. My question about Enums was basically if there is a way to define a list of possible values as well as string equivalents to those values. Adding the code_book field to the Config was how I got around it, but I'm thinking there must be a cleaner way.

[jeffzi]

Adding the code_book field to the Config was how I got around it,

The string equivalents are not used at all in your example. You want to validate that all elements of a column are in a set of int and string values?