stamtools/
split.rs

1use stam::*;
2use std::collections::BTreeSet;
3
4#[derive(Copy, Clone, Debug, PartialEq)]
5pub enum SplitMode {
6    Retain,
7    Delete,
8}
9
10pub fn split<'a>(
11    store: &'a mut AnnotationStore,
12    queries: Vec<Query<'a>>,
13    mode: SplitMode,
14    verbose: bool,
15) {
16    let mut split_annotations: BTreeSet<AnnotationHandle> = BTreeSet::new();
17    let mut split_resources: BTreeSet<TextResourceHandle> = BTreeSet::new();
18    let mut split_datasets: BTreeSet<AnnotationDataSetHandle> = BTreeSet::new();
19    for (i, query) in queries.into_iter().enumerate() {
20        eprintln!("Running query #{}...", i + 1);
21        let iter = store.query(query).expect("query failed");
22        for resultrow in iter {
23            if let Ok(result) = resultrow.get_by_name_or_last(Some("split")) {
24                match result {
25                    QueryResultItem::None => {}
26                    QueryResultItem::Annotation(annotation) => {
27                        split_annotations.insert(annotation.handle());
28                    }
29                    QueryResultItem::TextResource(resource) => {
30                        split_resources.insert(resource.handle());
31                    }
32                    QueryResultItem::AnnotationDataSet(dataset) => {
33                        split_datasets.insert(dataset.handle());
34                    }
35                    _ => {
36                        eprintln!("Error: Obtained result type is invalid for split, only ANNOTATION, RESOURCE and DATASET work.");
37                    }
38                }
39            }
40        }
41    }
42    eprintln!(
43        "Gathered {} annotations, {} resources, {} datasets",
44        split_annotations.len(),
45        split_resources.len(),
46        split_datasets.len()
47    );
48
49    match mode {
50        SplitMode::Delete => {
51            //split_* sets contains resources to delete
52            if !split_resources.is_empty() {
53                for resource in split_resources {
54                    if verbose {
55                        eprintln!("removing resource...");
56                    }
57                    let _ = store.remove_resource(resource); //we can just ignore handle error at this point
58                }
59            }
60            if !split_datasets.is_empty() {
61                for dataset in split_datasets {
62                    if verbose {
63                        eprintln!("removing dataset...");
64                    }
65                    let _ = store.remove_dataset(dataset); //we can just ignore handle error at this point
66                }
67            }
68            if !split_annotations.is_empty() {
69                for annotation in split_annotations {
70                    if verbose {
71                        eprintln!("removing annotation...");
72                    }
73                    let _ = store.remove_annotation(annotation); //we can just ignore handle error at this point
74                }
75            }
76        }
77        SplitMode::Retain => {
78            //split_* sets contains resources to retain
79            if !split_resources.is_empty() {
80                for resource in 0..store.resources_len() {
81                    let resource = TextResourceHandle::new(resource);
82                    if store.has(resource) && !split_resources.contains(&resource) {
83                        if verbose {
84                            eprintln!("removing resource #{}...", resource.as_usize());
85                        }
86                        let _ = store.remove_resource(resource); //we can just ignore handle error at this point
87                    }
88                }
89            }
90            if !split_datasets.is_empty() {
91                for dataset in 0..store.datasets_len() {
92                    let dataset = AnnotationDataSetHandle::new(dataset);
93                    if store.has(dataset) && !split_datasets.contains(&dataset) {
94                        if verbose {
95                            eprintln!("removing dataset #{}...", dataset.as_usize());
96                        }
97                        let _ = store.remove_dataset(dataset); //we can just ignore handle error at this point
98                    }
99                }
100            }
101            if !split_annotations.is_empty() {
102                for annotation in 0..store.annotations_len() {
103                    let annotation = AnnotationHandle::new(annotation);
104                    if store.has(annotation) && !split_annotations.contains(&annotation) {
105                        if verbose {
106                            eprintln!("removing annotation #{}...", annotation.as_usize());
107                        }
108                        let _ = store.remove_annotation(annotation); //we can just ignore handle error at this point
109                    }
110                }
111            }
112        }
113    }
114}
stamtools/split.rs

stamtools/
split.rs
