<commit_before>use std::collections::HashMap;
use std::cmp::{max, min};
use std::convert::AsRef;
use std::io;
use std::fs;
use std::path::Path;

use bio::io::gff::{self, GffType};
use itertools::{GroupBy, Group, Itertools};
use linked_hash_map::LinkedHashMap;

use {Coord, Exon, ExonFeatureKind as EFK, Gene, GBuilder, Strand, Transcript, Error};

// Various commonly-used feature column values
const GENE_STR: &'static str = "gene";
const TRANSCRIPT_STR: &'static str = "transcript";
const EXON_STR: &'static str = "exon";
const UTR_STR: &'static str = "UTR";
const UTR5_STR: &'static str = "UTR5";
const UTR3_STR: &'static str = "UTR3";
const CDS_STR: &'static str = "CDS";
const START_CODON_STR: &'static str = "start_codon";
const STOP_CODON_STR: &'static str = "stop_codon";

// Value for unknown columns.
const UNK_STR: &'static str = ".";

// Commonly-used attribute keys.
const GENE_ID_STR: &'static str = "gene_id";
const TRANSCRIPT_ID_STR: &'static str = "transcript_id";

// Initial coordinate for features.
const INIT_COORD: (u64, u64) = (::std::u64::MAX, ::std::u64::MIN);


pub struct Reader<R: io::Read> {
    inner: gff::Reader<R>,
}

impl<R: io::Read> Reader<R> {

    pub fn from_reader(in_reader: R, gff_type: GffType) -> Reader<R> {
        Reader {
            inner: gff::Reader::new(in_reader, gff_type),
        }
    }

    pub fn genes(&mut self) -> GffGenes<R> {
        GffGenes {
            inner: self.records().group_by(GffGenes::<R>::gene_groupf),
        }
    }

    fn records(&mut self) -> GffRecords<R> {
        GffRecords {
            inner: self.inner.records()
        }
    }
}

impl Reader<fs::File> {
    pub fn from_file<P: AsRef<Path>>(path: P, gff_type: GffType) -> io::Result<Self> {
        fs::File::open(path).map(|file| Reader::from_reader(file, gff_type))
    }
}

struct GffRecords<'a, R: 'a> where R: io::Read {
    inner: gff::Records<'a, R>,
}

impl<'a, R> Iterator for GffRecords<'a, R> where R: io::Read {

    type Item = Result<gff::Record, Error>;

    fn next(&mut self) -> Option<Result<gff::Record, Error>> {
        self.inner.next()
            .map(|row| {
                row.or_else(|err| Err(Error::from(err)))
            })
    }
}

pub struct GffGenes<'a, R: 'a> where R: io::Read, {
    inner: GroupBy<GxGroupKey, GffRecords<'a, R>, GxGroupFunc>,
}

type GxGroupKey = Option<(String, String, Strand)>;

type GxGroupFunc = fn(&Result<gff::Record, Error>) -> GxGroupKey;

type GxGroupedRecs<'a, 'b, R> = Group<'b, GxGroupKey, GffRecords<'a, R>, GxGroupFunc>;

type RawTrxCoord = (Coord<u64>, Vec<Coord<u64>>, Option<Coord<u64>>);

impl<'a, R> GffGenes<'a, R> where R: io::Read {

    fn gene_groupf(result: &Result<gff::Record, Error>) -> GxGroupKey {
        result.as_ref().ok()
            .map(|res| {
                let gene_id = res.attributes().get(GENE_ID_STR)
                    .map(|gid| gid.as_str())
                    .unwrap_or("");
                let seq_name = String::from(res.seqname());
                let strand = res.strand().unwrap_or(Strand::Unknown);

                (String::from(gene_id), seq_name, strand)<commit_msg>Add filter function before grouping in GTF genes() iterator<commit_after>use std::collections::HashMap;
use std::cmp::{max, min};
use std::convert::AsRef;
use std::io;
use std::iter::Filter;
use std::fs;
use std::path::Path;

use bio::io::gff::{self, GffType};
use itertools::{GroupBy, Group, Itertools};
use linked_hash_map::LinkedHashMap;

use {Coord, Exon, ExonFeatureKind as EFK, Gene, GBuilder, Strand, Transcript, Error};

// Various commonly-used feature column values
const GENE_STR: &'static str = "gene";
const TRANSCRIPT_STR: &'static str = "transcript";
const EXON_STR: &'static str = "exon";
const UTR_STR: &'static str = "UTR";
const UTR5_STR: &'static str = "UTR5";
const UTR3_STR: &'static str = "UTR3";
const CDS_STR: &'static str = "CDS";
const START_CODON_STR: &'static str = "start_codon";
const STOP_CODON_STR: &'static str = "stop_codon";

// Value for unknown columns.
const UNK_STR: &'static str = ".";

// Commonly-used attribute keys.
const GENE_ID_STR: &'static str = "gene_id";
const TRANSCRIPT_ID_STR: &'static str = "transcript_id";

// Initial coordinate for features.
const INIT_COORD: (u64, u64) = (::std::u64::MAX, ::std::u64::MIN);


pub struct Reader<R: io::Read> {
    inner: gff::Reader<R>,
}

impl<R: io::Read> Reader<R> {

    pub fn from_reader(in_reader: R, gff_type: GffType) -> Reader<R> {
        Reader {
            inner: gff::Reader::new(in_reader, gff_type),
        }
    }

    pub fn genes(&mut self) -> GffGenes<R> {
        GffGenes {
            inner: self.records()
                .filter(GffGenes::<R>::gene_filter_func as GxFilterFunc)
                .group_by(GffGenes::<R>::gene_group_func),
        }
    }

    fn records(&mut self) -> GffRecords<R> {
        GffRecords {
            inner: self.inner.records()
        }
    }
}

impl Reader<fs::File> {
    pub fn from_file<P: AsRef<Path>>(path: P, gff_type: GffType) -> io::Result<Self> {
        fs::File::open(path).map(|file| Reader::from_reader(file, gff_type))
    }
}

struct GffRecords<'a, R: 'a> where R: io::Read {
    inner: gff::Records<'a, R>,
}

impl<'a, R> Iterator for GffRecords<'a, R> where R: io::Read {

    type Item = Result<gff::Record, Error>;

    fn next(&mut self) -> Option<Result<gff::Record, Error>> {
        self.inner.next()
            .map(|row| {
                row.or_else(|err| Err(Error::from(err)))
            })
    }
}

pub struct GffGenes<'a, R: 'a> where R: io::Read, {
    inner: GroupBy<GxGroupKey, GxRecords<'a, R>, GxGroupFunc>,
}

type GxFilterFunc = fn(&Result<gff::Record, Error>) -> bool;

type GxRecords<'a, R> = Filter<GffRecords<'a, R>, GxFilterFunc>;

type GxGroupKey = Option<(String, String, Strand)>;

type GxGroupFunc = fn(&Result<gff::Record, Error>) -> GxGroupKey;

type GxGroupedRecs<'a, 'b, R> = Group<'b, GxGroupKey, GxRecords<'a, R>, GxGroupFunc>;

type RawTrxCoord = (Coord<u64>, Vec<Coord<u64>>, Option<Coord<u64>>);

impl<'a, R> GffGenes<'a, R> where R: io::Read {

    fn gene_filter_func(result: &Result<gff::Record, Error>) -> bool {
        match result {
            &Ok(ref rec) => rec.attributes().contains_key(GENE_ID_STR),
            &Err(_) => true,  // Err(_) is supposed to be handled elsewhere
        }
    }

    fn gene_group_func(result: &Result<gff::Record, Error>) -> GxGroupKey {
        result.as_ref().ok()
            .map(|res| {
                let gene_id = res.attributes().get(GENE_ID_STR)
                    .expect("'gene_id' attribute not found")
                    .clone();
                let seq_name = String::from(res.seqname());
                let strand = res.strand().unwrap_or(Strand::Unknown);

                (gene_id, seq_name, strand)
{'fraction_non_alphanumeric': 0.16926470588235293, 'fraction_numerical': 0.005294117647058823, 'mean_word_length': 3.3098859315589353, 'pattern_counts': {'":': 0, '<': 79, '<?xml version=': 0, '>': 97, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3066602', 'n_tokens_mistral': 2550, 'n_tokens_neox': 2488, 'n_words': 565}