variantplaner ¶

@classmethod
def minimal_schema(cls) -> dict[str, type]:
    """Get minimal schema of genotypes polars.LazyFrame."""
    return {
        "id": polars.UInt64,
    }

def from_path(self, path: pathlib.Path, /, **scan_csv_args: Unpack[ScanCsv]) -> int:
    """Fill object with file point by pathlib.Path.

    Argument:
    path: path of input file

    Returns: Number of contigs line view
    """
    csv = Csv()
    csv.from_path(path, **scan_csv_args)
    self.lf = csv.lf

    self.__compute_offset()

    return self.lf.collect(engine="cpu").shape[0]

def from_vcf_header(self, header: VcfHeader) -> int:
    """Fill a object with VcfHeader.

    Argument:
       header: VcfHeader

    Returns: Number of contigs line view
    """
    contigs_id = re.compile(r"ID=(?P<id>[^,]+)")
    contigs_len = re.compile(r"length=(?P<length>[^,>]+)")

    count = 0
    contigs2len: dict[str, list] = {"contig": [], "length": []}
    for contig_line in header.contigs:
        if (len_match := contigs_len.search(contig_line)) and (id_match := contigs_id.search(contig_line)):
            contigs2len["contig"].append(id_match.groupdict()["id"])
            contigs2len["length"].append(int(len_match.groupdict()["length"]))
        count += 1

    self.lf = polars.LazyFrame(contigs2len, schema={"contig": polars.String, "length": polars.UInt64})

    self.__compute_offset()

    return count

@classmethod
def minimal_schema(cls) -> dict[str, type]:
    """Get minimal schema of genotypes polars.LazyFrame."""
    return {
        "id": polars.UInt64,
        "sample": polars.String,
    }

def samples_names(self) -> list[str]:
    """Get list of sample name."""
    return self.lf.select("sample").unique("sample").collect(engine="cpu").get_column("sample").to_list()

def from_path(self, input_path: pathlib.Path) -> None:
    """Read a pedigree file in [polars.LazyFrame](https://pola-rs.github.io/polars/py-polars/html/reference/lazyframe/index.html).

    Args:
        input_path: Path to pedigree file.

    Returns:
        A [polars.LazyFrame](https://pola-rs.github.io/polars/py-polars/html/reference/lazyframe/index.html) that contains ped information ('family_id', 'personal_id', 'father_id', 'mother_id', 'sex', 'affected')
    """
    self.lf = polars.scan_csv(
        input_path,
        separator="\t",
        has_header=False,
        null_values=["None", "unknown"],
        new_columns=[
            "family_id",
            "personal_id",
            "father_id",
            "mother_id",
            "sex",
            "affected",
        ],
        schema_overrides=Pedigree.minimal_schema(),
    )

@classmethod
def minimal_schema(
    cls,
) -> collections.abc.Mapping[str, polars._typing.PolarsDataType]:
    """Get schema of variants polars.LazyFrame."""
    return {
        "family_id": polars.String,
        "personal_id": polars.String,
        "father_id": polars.String,
        "mother_id": polars.String,
        "sex": polars.String,
        "affected": polars.Boolean,
    }

def to_path(self, output_path: pathlib.Path) -> None:
    """Write pedigree [polars.LazyFrame](https://pola-rs.github.io/polars/py-polars/html/reference/lazyframe/index.html) in ped format.

    Warning: This function performs [polars.LazyFrame.collect][] before write csv, this can have a significant impact on memory usage

    Args:
        lf: LazyFrame contains pedigree information.
        output_path: Path where write pedigree information.

    Returns:
        None
    """
    self.lf.collect(engine="cpu").write_csv(output_path, include_header=False, separator="\t")

@classmethod
def minimal_schema(cls) -> dict[str, type]:
    """Get schema of variants polars.LazyFrame."""
    return {
        "id": polars.UInt64,
        "chr": polars.String,
        "pos": polars.UInt64,
        "ref": polars.String,
        "alt": polars.String,
    }

def add_genotypes(self, genotypes_lf: Genotypes) -> None:
    """Add genotypes information in vcf."""
    for sample in genotypes_lf.samples_names():
        geno2sample = (
            genotypes_lf.lf.filter(polars.col("sample") == sample)
            .rename(
                {col: f"{sample}_{col}" for col in genotypes_lf.lf.collect_schema().names()[2:]},
            )
            .drop("sample")
        )
        self.lf = self.lf.join(geno2sample, on="id", how="full", coalesce=True)

def annotations(self, select_info: set[str] | None = None) -> Annotations:
    """Get annotations of vcf."""
    lf = self.lf.with_columns(self.lf.header.info_parser(select_info))

    return lf.drop("chr", "pos", "ref", "alt", "format", "info")

def from_path(
    self,
    path: pathlib.Path,
    chr2len_path: pathlib.Path | None,
    behavior: VcfParsingBehavior = VcfParsingBehavior.NOTHING,
) -> None:
    """Populate Vcf object with vcf file."""
    with xopen.xopen(path) as fh:
        try:
            self.header.from_lines(fh)
        except NotVcfHeaderError as e:
            raise NotAVCFError(path) from e

    chr2len = ContigsLength()
    if chr2len_path is not None:
        if chr2len.from_path(chr2len_path) == 0 and chr2len.from_vcf_header(self.header) == 0:
            raise NoContigsLengthInformationError
    elif chr2len.from_vcf_header(self.header) == 0:
        raise NoContigsLengthInformationError

    self.lf = polars.scan_csv(
        path,
        separator="\t",
        comment_prefix="#",
        has_header=False,
        schema_overrides=Vcf.schema(),
        new_columns=list(Vcf.schema().keys()),
    )

    schema = self.lf.collect_schema()
    self.lf = self.lf.rename(dict(zip(schema.names(), self.header.column_name(schema.len()))))
    self.lf = self.lf.cast(Vcf.schema())  # type: ignore # noqa: PGH003  polars 1.0 typing stuff

    if behavior & VcfParsingBehavior.MANAGE_SV:
        self.lf = self.lf.with_columns(self.header.info_parser({"SVTYPE", "SVLEN"}))

    if behavior & VcfParsingBehavior.KEEP_STAR:
        self.lf = self.lf.filter(polars.col("alt") != "*")

    self.lf = normalization.add_variant_id(self.lf, chr2len.lf)

    if behavior & VcfParsingBehavior.MANAGE_SV:
        self.lf = self.lf.drop("SVTYPE", "SVLEN", strict=False)

def genotypes(self, format_str: str = "GT:AD:DP:GQ") -> Genotypes:
    """Get genotype of vcf."""
    schema = self.lf.collect_schema()

    if "format" not in schema.names():
        raise NoGenotypeError

    lf = self.lf.select([*schema.names()[schema.names().index("format") :]])
    schema = lf.collect_schema()

    # Clean bad variant
    lf = lf.filter(polars.col("format").str.starts_with(format_str)).select(*schema.names()[1:])

    # Found index of genotype value
    col_index = {
        key: index
        for (index, key) in enumerate(
            format_str.split(":"),
        )
    }

    # Pivot value
    genotypes = Genotypes()
    genotypes.lf = lf.unpivot(index=["id"]).with_columns(
        [
            polars.col("id"),
            polars.col("variable").alias("sample"),
            polars.col("value").str.split(":"),
        ],
    )

    # Split genotype column in sub value
    col2expr = self.header.format_parser()

    genotypes.lf = genotypes.lf.with_columns(
        [
            polars.col("value").list.get(index).pipe(function=col2expr[col], col_name=col)
            for col, index in col_index.items()
        ],
    )

    # Select intrusting column
    genotypes.lf = genotypes.lf.select(["id", "sample", *[col.lower() for col in col_index]])

    if "gt".upper() in col2expr:
        genotypes.lf = genotypes.lf.filter(polars.col("gt") != 0)

    return genotypes

@classmethod
def schema(cls) -> collections.abc.Mapping[str, polars._typing.PolarsDataType]:
    """Get schema of Vcf polars.LazyFrame."""
    return {
        "chr": polars.String,
        "pos": polars.UInt64,
        "vid": polars.String,
        "ref": polars.String,
        "alt": polars.String,
        "qual": polars.String,
        "filter": polars.String,
        "info": polars.String,
    }

def set_variants(self, variants: Variants) -> None:
    """Set variants of vcf."""
    self.lf = variants.lf

def variants(self) -> Variants:
    """Get variants of vcf."""
    return self.lf.select(Variants.minimal_schema())

def column_name(self, number_of_column: int = MINIMAL_COL_NUMBER) -> typing.Iterator[str]:
    """Get an iterator of correct column name.

    Returns: String iterator
    """
    base_col_name = ["chr", "pos", "vid", "ref", "alt", "qual", "filter", "info"]

    yield from base_col_name

    if number_of_column > MINIMAL_COL_NUMBER and (samples := self.samples_index):
        yield "format"
        yield from (sample for (sample, _) in samples.items())

def format_parser(
    self,
    select_format: set[str] | None = None,
) -> dict[str, typing.Callable[[polars.Expr, str], polars.Expr]]:
    """Generate a list of [polars.Expr](https://pola-rs.github.io/polars/py-polars/html/reference/expressions/index.html) to extract genotypes information.

    **Warning**: Float values can't be converted for the moment they are stored as String to keep information

    Args:
    header: Line of vcf header.
    input_path: Path to vcf file.
    select_format: List of target format field.

    Returns:
    A dict to link format id to pipeable function with Polars.Expr

    Raises:
    NotVcfHeaderError: If all line not start by '#CHR'
    """
    format_re = re.compile(
        "ID=(?P<id>[A-Za-z_][0-9A-Za-z_.]*),Number=(?P<number>[ARG0-9\\.]+),Type=(?P<type>Integer|Float|String|Character)",
    )

    expressions: dict[str, typing.Callable[[polars.Expr, str], polars.Expr]] = {}

    for line in self._header:
        if line.startswith("#CHROM"):
            return expressions

        if not line.startswith("##FORMAT"):
            continue

        if (search := format_re.search(line)) and (not select_format or search["id"] in select_format):
            name = search["id"]
            number = search["number"]
            format_type = search["type"]

            if name == "GT":
                expressions["GT"] = VcfHeader.__format_gt
                continue

            if number == "1":
                if format_type == "Integer":
                    expressions[name] = VcfHeader.__format_one_int
                elif format_type == "Float":  # noqa: SIM114 Float isn't already support but in future
                    expressions[name] = VcfHeader.__format_one_str
                elif format_type in {"String", "Character"}:
                    expressions[name] = VcfHeader.__format_one_str
                else:
                    pass  # Not reachable

            elif format_type == "Integer":
                expressions[name] = VcfHeader.__format_list_int
            elif format_type == "Float":  # noqa: SIM114 Float isn't already support but in future
                expressions[name] = VcfHeader.__format_list_str
            elif format_type in {"String", "Character"}:
                expressions[name] = VcfHeader.__format_list_str
            else:
                pass  # Not reachable

    raise NotVcfHeaderError

def from_files(self, path: pathlib.Path) -> None:
    """Populate VcfHeader object with content of only header file.

    Args:
    path: Path of file

    Returns:
    None
    """
    with open(path) as fh:
        for full_line in fh:
            line = full_line.strip()
            self._header.append(line)

def from_lines(self, lines: typing.Iterator[str]) -> None:
    """Extract all header information of vcf lines.

    Line between start of file and first line start with '#CHROM' or not start with '#'

    Args:
    lines: Iterator of line

    Returns: None

    Raises:
    NotAVcfHeader: If a line not starts with '#'
    NotAVcfHeader: If no line start by '#CHROM'
    """
    for full_line in lines:
        line = full_line.strip()

        if not line.startswith("#"):
            raise NotVcfHeaderError

        if line.startswith("#CHROM"):
            self._header.append(line)
            return

        self._header.append(line)

    raise NotVcfHeaderError

def info_parser(self, select_info: set[str] | None = None) -> list[polars.Expr]:
    """Generate a list of [polars.Expr](https://pola-rs.github.io/polars/py-polars/html/reference/expressions/index.html) to extract variants information.

    Args:
    header: Line of vcf header
    input_path: Path to vcf file.
    select_info: List of target info field

    Returns:
    List of [polars.Expr](https://pola-rs.github.io/polars/py-polars/html/reference/expressions/index.html) to parse info columns.

    Raises:
    NotVcfHeaderError: If all line not start by '#CHR'
    """
    info_re = re.compile(
        r"ID=(?P<id>([A-Za-z_][0-9A-Za-z_.]*|1000G)),Number=(?P<number>[ARG0-9\.]+),Type=(?P<type>Integer|Float|String|Character)",
    )

    expressions: list[polars.Expr] = []

    for line in self._header:
        if line.startswith("#CHROM"):
            return expressions

        if not line.startswith("##INFO"):
            continue

        if (search := info_re.search(line)) and (not select_info or search["id"] in select_info):
            regex = rf"{search['id']}=([^;]+);?"

            local_expr = polars.col("info").str.extract(regex, 1)

            if search["number"] == "1":
                if search["type"] == "Integer":
                    local_expr = local_expr.cast(polars.Int64)
                elif search["type"] == "Float":
                    local_expr = local_expr.cast(polars.Float64)
                elif search["type"] in {"String", "Character"}:
                    pass  # Not do anything on string or character
                else:
                    pass  # Not reachable

            else:
                local_expr = local_expr.str.split(",")
                if search["type"] == "Integer":
                    local_expr = local_expr.cast(polars.List(polars.Int64))
                elif search["type"] == "Float":
                    local_expr = local_expr.cast(polars.List(polars.Float64))
                elif search["type"] in {"String", "Character"}:
                    pass  # Not do anything on string or character
                else:
                    pass  # Not reachable

            expressions.append(local_expr.alias(search["id"]))

    raise NotVcfHeaderError