<commit_before>require_relative 'parsers/ast/node'

module BELParser
  # ASTFilter filters types of {BELParser::Parsers::AST::Node}.
  class ASTFilter
    def initialize(*types)
      @types = types
    end

    def each(ast_source)
      if block_given?
        ast_source.each do |(line_number, line, ast_results)|
          selected = filter(ast_results)
          yield([line_number, line, selected]) unless selected.empty?
        end
      else
        enum_for(:each, ast_source)
      end
    end

    def filter(ast_enum)
      ast_enum.select do |node|
        @types.include?(node.type)
      end
    end
  end
end<commit_msg>enable running ast_filter as a main script<commit_after>require_relative 'parsers/ast/node'

module BELParser
  # ASTFilter filters types of {BELParser::Parsers::AST::Node}.
  class ASTFilter
    def initialize(*types)
      @types = types
    end

    def each(ast_source)
      if block_given?
        ast_source.each do |(line_number, line, ast_results)|
          selected = filter(ast_results)
          yield([line_number, line, selected]) unless selected.empty?
        end
      else
        enum_for(:each, ast_source)
      end
    end

    def filter(ast_enum)
      ast_enum.select do |node|
        @types.include?(node.type)
      end
    end
  end
end

if __FILE__ == $PROGRAM_NAME
  require_relative 'ast_generator'

  types = ARGV.map(&:to_sym)
  ast   = BELParser::ASTGenerator.new.each($stdin)
  BELParser::ASTFilter.new(*types).each(ast) do |(line_number, line, results)|
    puts "#{line_number}: #{line}"
    results.each do |ast|
      puts ast.to_s(1)
    end
  end
end
{'fraction_non_alphanumeric': 0.10122699386503067, 'fraction_numerical': 0.0006134969325153375, 'mean_word_length': 2.775462962962963, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '4864776', 'n_tokens_mistral': 535, 'n_tokens_neox': 522, 'n_words': 126}