ps expand recursive example

example.nu

use std *

# Convert a value into a list.
export def 'into list' []: any -> list<any> {
  append null
}


# { type: <'list'>     , length   : <int>, values: list<detailed_type>     }
# { type: <'record'>   , columns  : record<any, detailed_type>             }
# { type: <'closure'>  , signature: { name: <string>, category: <string> } }
# { type: <'error'>    , subtype  : <string>                               }
# { type: <'binary'>   , length   : <int>                                  }
# { type: <'cell-path'>, length   : <int>                                  }
# { type: <'custom'>   , subtype  : <string>                               }

# Expects an input that is already described-- digs into the description record and gets the appropriate output
export def get-nested-type [
  --keep-streams(-k)
] {
  match $in {
    { type: 'stream', subtype: $subtype } => { if $keep_streams { 'stream' } else { $subtype | get-nested-type --keep-streams=($keep_streams) } }
    { type: 'stream' } => { if $keep_streams { 'stream' } else { error make { msg: $'No subtype for this stream?' } } }
    { type: $type } => { $type }
    $type => { $type }
  }
}

# Gets the type of an input. Returns one of:
# - binary
# - bool
# - cell-path
# - closure
# - datetime
# - string (directory, path)
# - duration
# - filesize
# - float
# - glob
# - int
# - list
# - nothing
# - number
# - range
# - record
export def type [
  --keep-streams(-k) # If the input is a stream, should `stream` be returned, or should `.subtype.type` be returned?
  --no-collect(-n) # Do not collect the input
] {
  describe --detailed --no-collect=($no_collect) | get-nested-type --keep-streams=($keep_streams)
}

export def create-column-closure [
  --ignore-errors(-i) # Whether or not to ignore errors
  --default(-d): string = 'null' # The default value to return
  --allow-invalid-keys(-k) # If true, invalid keys (anything other than `binary`, `int`, ...) will be returned instead of defaulted.
  input: any
  # The method to get the column. Can either be:
  # - closure(any) -> any
  # - cell-path
  # - string (converts into cell-path)
  # - int (converts into cell-path)
  # - list<any> (converts into cell-path)
  # - list<record<value: any, optional: bool>> (converts into cell-path)
]: [ any -> closure ] {

  let into_valid_key = {|key|
    if $allow_invalid_keys {
      $key
    } else (match ($key | type) {
      'binary' |
      'int' |
      'number' |
      'string' |
      'glob' |
      'bool' |
      'filesize' |
      'date' |
      'datetime' |
      'duration' => { $key }
      _ => { $default }
    })
  }

  match ($input | type) {
    'closure' => { {|value: any| $value | do --ignore-errors=($ignore_errors) $input | do $into_valid_key $in } }
    'cell-path' => { {|value: any| $value | get --ignore-errors=($ignore_errors) $input | do $into_valid_key $in } }
    'string' => {
      $input | split row '.' | each {|part|
        let trimmed = $part | str replace -r '\?$' ''
        { value: $trimmed, optional: ($trimmed != $part) }
      } | into cell-path | create-column-closure $in --ignore-errors=($ignore_errors) --default=($default) --allow-invalid-keys=($allow_invalid_keys)
    }
    'int' => {
      $input | into cell-path | create-column-closure $in --ignore-errors=($ignore_errors) --default=($default) --allow-invalid-keys=($allow_invalid_keys)
    }
    'list' => {
      $input | into cell-path | create-column-closure $in --ignore-errors=($ignore_errors) --default=($default) --allow-invalid-keys=($allow_invalid_keys)
    }
    $x => {
      error make { msg: $'Unsupported input type: ($x)' }
    }
  }
}

# creates a map using the specified key
export def create-map [
  --ignore-errors(-i) # Whether or not to ignore errors
  --default(-d): string = 'null' # The default value to return
  --group(-g)
  # Whether or not to group the results -- if `true`, then the output will look like: { [key: any]: list<any> }
  # if `false`, then the output will look like: { [key: any]: any }
  --to-table(-t) # If true, the output will be a `table<key: any, value: any>`
  column: any
  # The method to get the column. Can either be:
  # - closure(any) -> any
  # - cell-path
  # - string (converts into cell-path)
  # - int (converts into cell-path)
  # - list<any> (converts into cell-path)
  # - list<record<value: any, optional: bool>> (converts into cell-path)
]: [
  list -> record
  list -> table
  table -> record
  table -> table
] {
  let input = $in | into list

  let column_closure = create-column-closure $column --ignore-errors=($ignore_errors) --default=($default) --allow-invalid-keys=($to_table)

  let mapped_values = $input | each {|item|
    { key: ($item | do $column_closure $item | into string), value: $item }
  }

  let outputs = if ($group) {
    if $to_table {
      $mapped_values | group-by --to-table key | rename -c { group: 'key', items: 'value' } | upsert value { get value }
    } else {
      $mapped_values | group-by key --to-table | upsert items { get value } | transpose -rid
    }
  } else {
    if $to_table {
      $mapped_values | transpose -rid | transpose key value
    } else {
      $mapped_values | transpose -rid
    }
  }

  $outputs
}

export def 'ps expand' [] {
  let processes = ps --long
    | reject -i env_diff priority user_sid
    | upsert pid { into string }
    | upsert ppid { into string }

  $env.acc = {
    raw: {
      processes : $processes,
      pid_map   : ($processes | create-map -i pid),
      ppid_map  : ($processes | create-map -ig ppid),
    },
    expanded_map: {}
    parent_map: {}
  }

  def --env expand-process [] {
    let process  = $in
    let expanded = $env.acc.expanded_map | get -i $process.pid
    if ($expanded | is-not-empty) {
      return $expanded
    }

    # Now we really have to expand it. Get the children.
    let children = $env.acc.raw.ppid_map | get -i $process.pid
    if ($children | is-empty) {
      let expanded = $process | upsert children { [] }
      $env.acc.expanded_map = (
        $env.acc.expanded_map
        | upsert $process.pid $expanded
      )
      return $expanded
    }


    # We have children to expand, so expand the children.
    mut expanded_children = []
    for $child in $children {
      $expanded_children ++= ($child | expand-process)
    }

    let expanded = $process | upsert children $expanded_children
    $env.acc.expanded_map = (
      $env.acc.expanded_map
      | upsert $process.pid $expanded
    )

    $expanded
  }

  def --env expand-process-parent [] {
    let process  = $in
    let expanded = $env.acc.parent_map | get -i $process.pid
    if ($expanded | is-not-empty) {
      return $expanded
    }

    let parent = $env.acc.expanded_map | get -i $process.ppid
    if ($parent | is-empty) {
      let output = $process | upsert parent { null }
      $env.acc.parent_map = ( $env.acc.parent_map | upsert $process.pid $output )
      return $output
    }

    let expanded_parent = $parent | expand-process-parent
    let expanded_proc   = $process | upsert parent $expanded_parent
    $env.acc.parent_map = ($env.acc.parent_map | upsert $expanded_proc.pid $expanded_proc)
    $expanded_proc
  }

  for $proc in $env.acc.raw.processes {
    let expanded = $proc | expand-process
    $env.acc.expanded_map = ($env.acc.expanded_map | upsert $expanded.pid $expanded)
  }

  for $proc in ($env.acc.expanded_map | values) {
    let expanded = $proc | expand-process-parent
    $env.acc.parent_map = ($env.acc.parent_map | upsert $expanded.pid $expanded)
  }

  $env.acc.parent_map
  | values
  | move parent children --before pid
}