GemTree

gemtree.rb

#!/usr/bin/env ruby
#  -*- coding: utf-8 -*-

require 'optparse'
require 'ostruct'

# using list from graph gem
COLOR_SCHEME_MAX = {
  accent:   8,  blues: 9,  brbg:   11, bugn:    9,
  dark2:    8,  gnbu:  9,  greens: 9,  greys:   9,
  oranges:  9,  orrd:  9,  paired: 12, pastel1: 9,
  pastel2:  8,  piyg:  11, prgn:   11, pubu:    9,
  pubugn:   9,  puor:  11, purd:   9,  purples: 9,
  rdbu:     11, rdgy:  11, rdylbu: 11, rdylgn:  11,
  reds:     9,  set1:  9,  set2:   8,  set3:    12,
  spectral: 11, ylgn:  9,  ylgnbu: 9,  ylorbr:  9,
  ylorrd:   9
}

# :nodoc:
class Array
  def to_h
    Hash[self]
  end unless method_defined?(:to_h)
end

# Helper class for loading a gemfile with instance_eval
class GemFileLoader
  attr_reader :groups

  attr_reader :specs

  def initialize(file, fgroups = [])
    @groups = Hash.new { |hash, key| hash[key] = [] }
    @group, @specs = :default, []

    # Change the dir to the Gemfile itself,
    # because the paths to the gemspecs are relative
    Dir.chdir(File.dirname(file)) do
      instance_eval(File.read(File.basename(file)))
    end
    filter_groups(fgroups) unless fgroups.empty?
  end

  def ruby(*)
  end

  def gemspec(path: '.', name: nil, development_group: :development, **)
    # change the path to the dir of the gemspec,
    # because it might use git inside and otherwise it cant find anything
    Dir.chdir(File.dirname(path)) do
      # use default name with is based on the current dir
      name ||= File.basename(Dir.pwd)
      name += '.gemspec'

      s = eval(File.read(name))
      # add spec to extra spec list,
      # add them later to the all specs Hashs so it can be found
      @specs << s
      # add the spec to the gurrent gem group
      gem s.name
      # add the development stuff into extra development_group
      s.development_dependencies.each { |dep| _add(dep, development_group) }
    end
  end

  def method_missing(m, *)
    super unless [:source, :git, :path, :platforms].include?(m)
    yield if block_given?
  end

  def gem(name, *versions, group: nil, groups: @group, **)
    group ||= groups
    _add(Gem::Dependency.new(name, *versions), group)
  end

  def group(*name)
    new_group, @group = @group, name
    yield
    @group = new_group
  end

  def all_names
    @specs.map(&:name) | @groups.each_value.flat_map { |v| v.map(&:name) }
  end

  def filter_groups(gr)
    @groups.keep_if { |k, _| gr.include?(k.to_s) }
  end

  private

  def _add(dep, group = @group)
    Array(group).each { |g| @groups[g] << dep }
  end
end

options = OpenStruct.new
options.prerelease = false
options.develop = false
options.missing = true
options.version = false

options.gems = []

options.sort = 1
options.grouping = true

options.output = '-'
options.encoding = 'utf-8'

options.gemfile = nil
options.only_gemfile = false
options.gemfile_groups = []
options.gemfile_colorscheme = nil # "pastel1"

options.splines = nil

options.both_directions = false

options.develop_color = 'green'

options.unnecessary_color = 'blue'

options.missing_shape = 'hexagon'
options.missing_color = 'red'
options.missing_dev_color = 'orange'

options.gemfile_group_shape = 'box3d'

options.gemfile_color = 'yellow'

options.gemfile_group_style = 'filled'
options.gemfile_style = 'wedged' # 'striped'

options.update_style = 'dashed'
options.default_style = 'dotted'

opt_parser = OptionParser.new do |opts|
  opts.banner = "Usage: #{$PROGRAM_NAME} [options]"

  opts.separator ''
  opts.separator 'Specific options:'

  opts.on(
    '-o', '--output path',
    'output file where the stuff is written too (default is STDOUT)'
  ) do |o|
    options.output = o
  end

  opts.on(
    '-g', '--gems names', Array,
    'names of gems that should be added as graph root'
  ) do |gems|
    options.gems.concat(gems)
  end

  opts.on(
    '--gemfile gemfile', 'path option for optional gemfile'
  ) do |g|
    options.gemfile = g
  end

  opts.on(
    '--gemgroup names', Array,
    'names of groups of the gemfile that should be used'
  ) do |gems|
    options.gemfile_groups.concat(gems)
  end

  opts.on(
    '-s', '--sort order', 'how the nodes should be sorted, ' \
    '1 = by size, 0 = unsorted, -1 = negative by size'
  ) do |s|
    options.sort = s.to_i
  end

  #  Boolean switch.
  opts.on('-d', '--[no-]develop', 'Show Developdepency') do |d|
    options.develop = d
  end
  opts.on('-p', '--[no-]prerelease', 'Enables prerelease gems') do |p|
    options.prerelease = p
  end

  opts.on('-m', '--[no-]missing', 'Show missing') do |m|
    options.missing = m
  end
  opts.on('--[no-]grouping', 'Group parent nodes together') do |g|
    options.grouping = g
  end

  opts.on(
    '--[no-]only-gemfile',
    'Shows only gems which are part of the Gemfile'
  ) do |og|
    options.only_gemfile = og
  end

  opts.on('-V', '--[no-]verbose', 'Show always version') do |v|
    options.version = v
  end

  opts.separator 'Graph options:'

  opts.on('--splines spline', 'graph option for splines') do |s|
    options.splines = s
  end

  opts.on(
    '-b', '--[no-]both-directions',
    'Draw the connections in both directions if possible'
  ) do |b|
    options.both_directions = b
  end

  opts.on(
    '--gemfile-colorscheme scheme',
    'name of the colorscheme that is used'
  ) do |gcs|
    options.gemfile_colorscheme = gcs
  end

  opts.separator 'General:'

  #  No argument, shows at tail.  This will print an options summary.
  #  Try it and see!
  opts.on_tail('-h', '--help', 'Show this message') do
    puts opts
    exit
  end
end

opt_parser.parse!(ARGV)

# :nodoc:
class GemFilter
  def initialize(options)
    # Create Hashs for specs, deps and dev_deps
    name, specs, dep, dev_dep = data(options.prerelease)

    # combine the Arrays into Hashs
    set(name.zip(specs).to_h, name.zip(dep).to_h, name.zip(dev_dep).to_h)
    @develop = options.develop
  end

  def data(prerelease)
    Gem::Specification.latest_specs(prerelease).map do |so|
      [so.name, so, so.runtime_dependencies, so.development_dependencies]
    end.transpose
  end

  def set(specs, dep, dev_dep)
    @specs, @dep, @dev_dep = specs, dep, dev_dep
  end

  def new_specs
    [@specs, @dep, @dev_dep]
  end

  def add_spec(*specs)
    specs.each do |spec|
      @specs[spec.name] = spec
      @dep[spec.name] = spec.runtime_dependencies
      @dev_dep[spec.name] = spec.development_dependencies
    end
  end

  # helper function for filter gems move from source to sink using filter
  # using fnmatch for matching
  # return new values
  def filter_select(source, filter, sink)
    f = source.select { |s| filter.any? { |m| File.fnmatch(m, s) } }
    sink.merge! f
    f.values.flatten.flat_map(&:name).uniq
  end

  # does filter the specs, deps and dev deps for only the giving gems
  def filter_gems(add_gems = [], new_specs = {}, new_dep = {}, new_dev_dep = {})
    # filter only adding gems from specs into new_specs
    filter_select(@specs, add_gems, new_specs)

    # do the same with the dependencies, remember the ones that got added
    adding = filter_select(@dep, add_gems, new_dep)

    # for development build add also the
    # development dependencies of the new gems
    adding.concat filter_select(@dev_dep, add_gems, new_dev_dep) if @develop

    # remove already added gems from the adding list
    # to prevent circular dependencies
    adding.uniq!
    adding.delete_if { |n| new_specs.include?(n) }

    # if there are still some new gems, repeat the process
    filter_gems(adding, new_specs, new_dep, new_dev_dep) unless adding.empty?

    # return the new specs, deps and dev deps
    [new_specs, new_dep, new_dev_dep]
  end

  def filter_gems!(add_gems = [])
    set(*filter_gems(add_gems))
  end
end

# :nodoc:
class GemTree
  def initialize(options)
    @options = options

    filter = GemFilter.new(options)

    # load a gem file if it's given in the options
    init_gemfileloader(filter) if options.gemfile

    # filter gemspecs for only wanted gems
    filter.filter_gems!(options.gems) unless options.gems.empty?

    @specs, @dep, @dev_dep = filter.new_specs

    # use an Hash with default proc to get a list
    # of all specs for a gem on command
    @all_specs = Hash.new do |hash, key|
      hash[key] = Gem::Specification.find_all_by_name(key)
    end

    @checker = Checker.new(@specs, @all_specs, options)
  end

  def init_gemfileloader(filter)
    # parse gemfile using instance_eval and capture the output
    # also remove unwanted groups
    @loader = GemFileLoader.new(@options.gemfile, @options.gemfile_groups)

    filter.add_spec(*@loader.specs)

    # with this remove all gems that are not part of the gemfile
    filter.filter_gems!(@loader.all_names) if @options.only_gemfile

    @colorscheme = @options.gemfile_colorscheme = validate_colorscheme_name(
      @options.gemfile_colorscheme, @loader.groups.size
    ) if @options.gemfile_colorscheme
  end

  def validate_colorscheme_name(name, val)
    sym_name = name.to_sym
    scheme_max = COLOR_SCHEME_MAX
    return if scheme_max[sym_name].nil? || scheme_max[sym_name] < val
    name + [val, 3].max.to_s
    # s, m = COLOR_SCHEME_MAX.find do |k, v|
    #   /\A#{k}(?<n>\d+)\Z/ =~ name && n.to_i.between?(3, v)
    # end
  end

  # :nodoc:
  class Checker
    def initialize(specs, all_specs, options)
      @specs, @all_specs, @options = specs, all_specs, options
    end
    # add version label but only if it should show the version
    # and it does not require the latest one
    def add_version(par, opt)
      return if par.latest_version?
      opt[:xlabel] = "\"#{par.requirements_list.join('\n')}\""
    end

    def add_missing(par, opt, extra_nodes)
      opt[:color] = @options.send(
        par.type == :runtime ? :missing_color : :missing_dev_color
      )

      # add extra options to make the node appear missing too
      extra_nodes[par.name].update(
        shape: @options.missing_shape,
        color: opt[:color]
      ) { |_, v, _| v }
    end

    def find_dep(gem, enum)
      d = enum.find { |dep| dep.name == gem }
      yield d if d && block_given?
      d
    end

    def update_check_unnecessary(gem, spec, dep, opt)
      # for runtime check if requirement is unnecessary
      if dep && dep[gem].any? do |m|
        Array(dep[m.name]).any? { |d| d.name == spec.name }
      end
        return opt[:color] = @options.unnecessary_color
      end
      false
    end

    def update_check_runtime(gem, spec, dep, opt)
      return unless gem
      if @options.both_directions
        if find_dep(gem, spec.runtime_dependencies)
          opt[[gem, s.name].min == gem ? :dir : :already] = 'both'
        end
        # if one of them is development, do it elsewhere
        if @options.develop
          opt[:already] = find_dep(gem, spec.development_dependencies)
        end
      end

      # for runtime check if requirement is unnecessary
      update_check_unnecessary(gem, spec, dep, opt)
    end

    def update_check_development_both_runtime(gem, s, satisfy, opt, new_color)
      find_dep(gem, s.runtime_dependencies) do |rd|
        opt[:dir] = 'both'
        new_color << (satisfy.call(rd) ? 'black' : @options.missing_color)
      end
    end

    def update_check_development_both_devel(gem, s, satisfy, opt, new_color)
      find_dep(gem, s.development_dependencies) do |dd|
        opt[[gem, s.name].min == gem ? :dir : :already] = 'both'
        new_color << @options.missing_dev_color unless satisfy.call(dd)
      end
    end

    def update_check_development_both(gem, s, opt)
      satisfy, new_color = @specs[gem].method(:satisfies_requirement?), []

      update_check_development_both_runtime(gem, s, satisfy, opt, new_color)
      update_check_development_both_devel(gem, s, satisfy, opt, new_color)

      opt.update(color: new_color) { |_, o, n| [*o, *n].uniq }
    end

    def update_check_development(gem, s, _, opt)
      # if dep is not in parameters, it's a develop dependency
      opt[:color] = @options.develop_color

      return unless gem && s.name != gem && @options.both_directions
      update_check_development_both(gem, s, opt)
      false
    end

    def satisfy_check(spec, par, opt)
      return true if spec.satisfies_requirement? par
      par_list = @all_specs[par.name]
      if par_list.empty? || par_list.none? { |s| s.satisfies_requirement? par }
        # none of the installed versions does satisfies requirement
        # treat it as missing too
        return false
      else
        # one does satisfies, but it's not the newest version,
        # parent gem does need to be updated
        opt[:style] = @options.update_style
      end
    end

    def update_check_extra(missing, version, par, opt, extra_nodes)
      # add styles if it's missing and the missing flag is given
      add_missing(par, opt, extra_nodes) if missing && @options.missing

      add_version(par, opt) if version
    end

    # checks if gem is installed and if requirement is unnecessary
    def update_check(gem, par, dep, extra_nodes)
      missing, opt, s = true, {}, @specs[par.name]

      version = @options.version || @all_specs[par.name].size > 1

      # check if parent is in specs, if not it's missing
      if s
        missing = false
        version = true if send("update_check_#{par.type}", gem, s, dep, opt)

        # check if none does satisfies
        version = missing = true unless satisfy_check(s, par, opt)
      end
      update_check_extra(missing, version, par, opt, extra_nodes)

      [missing, opt]
    end
  end

  def style_dep(dep, extra_nodes)
    # sort the nodes, for better reading of the dot file,
    # has slightly difference in the image output
    dep.sort_by { |_, v| v.size * @options.sort }.map do |gem, parents|
      # group parents by their options to put them together and short the file
      [gem, parents.group_by do |par|
        missing, opt = @checker.update_check(gem, par, dep, extra_nodes)

        next if opt[:already] || (missing && !@options.missing)

        opt
      end]
    end.to_h
  end

  def style_loader_style(d, extra_nodes)
    if @colorscheme && extra_nodes[d.name][:fillcolor]
      @options.gemfile_style
    else
      @options.gemfile_group_style
    end
  end

  def style_loader_group(g, list, color, extra_nodes)
    # group the list of dependencies of their options
    [g, list.group_by do |d|
      # style gem nodes with extra style hash
      extra_nodes[d.name].update(
        fillcolor: color, style: style_loader_style(d, extra_nodes)
      ) { |_, o, n| [*o, *n].uniq }

      missing, opt = @checker.update_check(nil, d, nil, extra_nodes)
      next if missing && !@options.missing

      opt
    end]
  end

  def style_loader(extra_nodes)
    @loader.groups.each.with_index(1).map do |(g, list), i|
      color = @colorscheme ? i : @options.gemfile_color
      # style the gemfile groups with extra style hash
      extra_nodes[g].update(
        shape:     @options.gemfile_group_shape,
        fillcolor: color,
        style:     @options.gemfile_group_style
      )
      style_loader_group(g, list, color, extra_nodes)
    end.to_h
  end

  def style_extra_node_update(new_styles, spec, spec_dep_gems, list)
    new_styles << @options.update_style if list.size > 1 && list.any? do |s|
      s.version != spec.version && (s.dependent_gems - spec_dep_gems).empty?
    end
  end

  def style_extra_node(opt, spec, spec_dep_gems, list)
    new_styles = []

    # mark gems that can be cleaned up
    style_extra_node_update(new_styles, spec, spec_dep_gems, list)

    # mark default gems
    new_styles << @options.default_style if list.any?(&:default_gem?)

    return if new_styles.empty?
    opt.update(style: new_styles) { |_, o, n| [*o, *n].uniq }
  end

  def style_extra_nodes(extra_nodes)
    # some addional extra options
    @specs.each do |gem, spec|
      spec_dep_gems = spec.dependent_gems
      next if spec_dep_gems.empty? && spec.dependencies.empty?
      style_extra_node(extra_nodes[gem], spec, spec_dep_gems, @all_specs[gem])
    end
  end

  def style_all(extra_nodes)
    output = { dep: style_dep(@dep, extra_nodes) }
    # if develop is enabled show also it's development dependency
    output[:dev_dep] = style_dep(@dev_dep, extra_nodes) if @options.develop
    output[:loader] = style_loader(extra_nodes) if @loader

    style_extra_nodes(extra_nodes)

    output
  end
end

def escape_list(list, filter = '%s', border = '', join = ',')
  list, border = Array(list), Array(border)
  return list.first if list.size == 1
  border[0] + list.map(&filter.method(:%)).join(join) + border[-1]
end

# helper function to escape the options for dot in a [a=b,c=d] format
def escape_opt(opt)
  return '' if opt.empty?
  '[' + opt.map do |k, v|
    "#{k}=" + escape_list(v, '%s', '"', k['color'] ? ':' : ',').to_s
  end.join(', ') + ']'
end

# helper function to escape the nodes
# if there is more than one for "a" or {"a","b"}
def escape_nodes(nodes)
  nodes = Array(nodes)
  nodes.size > 1 ? escape_list(nodes, '"%s"', %w({ })) : "\"#{nodes.first}\""
end

# helper function to print a node to nodes connection,
# print them as group if enabled
def put_connect(file, node, other, group = true, **opt)
  if group
    file.puts "\"#{node}\" -> #{escape_nodes(other)}#{escape_opt(opt)};"
  else
    Array(other).each { |o| put_connect(file, node, o, opt) }
  end
end

# puts options for a node. used only for nodes with extra options
# nods like that cant be grouped
def put_node(file, node, opt)
  file.puts "\"#{node}\"#{escape_opt(opt)};"
end

def write_groups(file, gem, options, groups)
  groups.each do |opt, parents|
    # print the connections, does group them together if wanted and possible
    next unless opt
    put_connect(file, gem, parents.map(&:name), options.grouping, opt)
  end
end

def write_extra_nodes(file, extra_nodes)
  # print the extra nodes, but only if they really have options
  extra_nodes.each do |gem, opt|
    put_node(file, gem, opt) unless opt.empty?
  end
end

def write_file_head(file, options)
  # adding extra options
  file.puts "splines=#{options.splines};" if options.splines

  cs = options.gemfile_colorscheme
  return unless options.gemfile && cs
  opt = {}

  opt[:shape] = 'rectangle' if options.gemfile_style == 'striped'
  opt[:colorscheme] = cs

  file.puts "node #{escape_opt(opt)};" unless opt.empty?
end

def write_file(file, sections, options, extra_nodes)
  file.puts 'digraph inheritance {'

  write_file_head(file, options)

  sections.each_value do |hash|
    hash.each { |gem, group| write_groups(file, gem, options, group) } if hash
  end

  write_extra_nodes(file, extra_nodes)

  file.puts '}'
end

# extra nodes flag to add options at runtime
extra_nodes = Hash.new { |hash, key| hash[key] = {} }

tree = GemTree.new(options)

sections = tree.style_all(extra_nodes)

if options.output == '-'
  write_file(STDOUT, sections, options, extra_nodes)
else
  File.open(options.output, 'w', encoding: options.encoding) do |file|
    write_file(file, sections, options, extra_nodes)
  end
end