Module: Docscribe::Infer::Returns

Defined in:
lib/docscribe/infer/returns.rb

Overview

Return type inference and rescue-conditional return extraction.

Constant Summary collapse

LAST_EXPR_TYPE_HANDLERS = 
{
  begin: :handle_begin_node,
  if: :handle_if_node,
  case: :handle_case_node,
  return: :handle_return_node,
  block: :handle_block_node,
  send: :handle_send_node
}.freeze

Class Method Summary collapse

Class Method Details

.assignment_name_and_value(node) ⇒ Array<(String, Parser::AST::Node)>

Note:

module_function: when included, also defines #assignment_name_and_value (instance visibility: private)

Extract the variable name and value expression from an assignment node.

Parameters:

  • node (Parser::AST::Node) —

    an assignment AST node (:lvasgn, :gvasgn, :ivasgn, :casgn)

Returns:

  • (Array<(String, Parser::AST::Node)>) —

    pair of variable name and value node



203
204
205
206
207
208
209
210
211
212
 
# File 'lib/docscribe/infer/returns.rb', line 203

def assignment_name_and_value(node)
  case node.type
  when :lvasgn, :gvasgn, :ivasgn
    [node.children[0].to_s, node.children[1]]
  when :casgn
    [node.children[0].to_s, node.children[2]]
  else
    [nil, nil]
  end
end

.build_local_variable_types(node) ⇒ Hash?

Note:

module_function: when included, also defines #build_local_variable_types (instance visibility: private)

Build a map of local/global/ivar/constant assignments to inferred types.

Parameters:

  • node (Parser::AST::Node) —

    AST node to walk

Returns:

  • (Hash, nil) 


176
177
178
179
180
181
182
 
# File 'lib/docscribe/infer/returns.rb', line 176

def build_local_variable_types(node)
  types = {} #: Hash[String, String]
  ASTWalk.walk(node) do |n|
    collect_assignment_type(n, types)
  end
  types.empty? ? nil : types
end

.collect_assignment_type(node, types) ⇒ void

Note:

module_function: when included, also defines #collect_assignment_type (instance visibility: private)

This method returns an undefined value.

Infer the type of a single assignment node and store it in the types hash.

Parameters:

  • node (Parser::AST::Node) —

    an assignment AST node

  • types (Hash) —

    the accumulated local variable type map



190
191
192
193
194
195
196
 
# File 'lib/docscribe/infer/returns.rb', line 190

def collect_assignment_type(node, types)
  name, value = assignment_name_and_value(node)
  return unless name && value

  inferred = Literals.type_from_literal(value, fallback_type: FALLBACK_TYPE)
  types[name] = inferred if inferred && inferred != FALLBACK_TYPE
end

.extract_def_body(node) ⇒ Parser::AST::Node?

Note:

module_function: when included, also defines #extract_def_body (instance visibility: private)

Extract the body child node from a ‘:def` or `:defs` AST node.

Parameters:

  • node (Parser::AST::Node) —

    a ‘:def` or `:defs` AST node

Returns:

  • (Parser::AST::Node, nil) 


101
102
103
104
105
106
 
# File 'lib/docscribe/infer/returns.rb', line 101

def extract_def_body(node)
  case node.type
  when :def then node.children[2]
  when :defs then node.children[3]
  end
end

.handle_begin_node(node, **opts) ⇒ Object

Note:

module_function: when included, also defines #handle_begin_node (instance visibility: private)

Handle ‘:begin` node for last_expr_type.

Parameters:

  • node (Object)
  • opts (Hash)

Returns:

  • (Object) 


220
221
222
 
# File 'lib/docscribe/infer/returns.rb', line 220

def handle_begin_node(node, **opts)
  run_last_expr_type(node.children.last, **opts)
end

.handle_block_node(node, **opts) ⇒ Object

Note:

module_function: when included, also defines #handle_block_node (instance visibility: private)

Handle ‘:block` node for last_expr_type.

Parameters:

  • node (Object)
  • opts (Hash)

Returns:

  • (Object) 


277
278
279
280
281
282
283
284
285
286
287
288
 
# File 'lib/docscribe/infer/returns.rb', line 277

def handle_block_node(node, **opts)
  send_node = node.children[0]
  if send_node&.type == :send
    recv = send_node.children[0]
    meth = send_node.children[1]
    rbs_type = resolve_rbs_for_send(recv, meth, opts[:core_rbs_provider], opts[:local_var_types],
                                    opts[:param_types])
    return rbs_type if rbs_type
  end

  run_last_expr_type(node.children[2], **opts)
end

.handle_case_node(node, **opts) ⇒ Object

Note:

module_function: when included, also defines #handle_case_node (instance visibility: private)

Handle ‘:case` node for last_expr_type.

Parameters:

  • node (Object)
  • opts (Hash)

Returns:

  • (Object) 


243
244
245
246
247
248
249
250
251
252
253
 
# File 'lib/docscribe/infer/returns.rb', line 243

def handle_case_node(node, **opts)
  branches = process_case_branches(node, **opts)
  if branches.empty?
    opts[:fallback_type]
  else
    branches.reduce do |a, b|
      unify_types(a, b, fallback_type: opts[:fallback_type] || 'untyped',
                        nil_as_optional: opts.fetch(:nil_as_optional, true))
    end
  end
end

.handle_if_node(node, **opts) ⇒ Object

Note:

module_function: when included, also defines #handle_if_node (instance visibility: private)

Handle ‘:if` node for last_expr_type.

Parameters:

  • node (Object)
  • opts (Hash)

Returns:

  • (Object) 


230
231
232
233
234
235
 
# File 'lib/docscribe/infer/returns.rb', line 230

def handle_if_node(node, **opts)
  t = run_last_expr_type(node.children[1], **opts)
  e = run_last_expr_type(node.children[2], **opts)
  unify_types(t, e, fallback_type: opts[:fallback_type] || 'untyped',
                    nil_as_optional: opts.fetch(:nil_as_optional, true))
end

.handle_return_node(node, **opts) ⇒ String?

Note:

module_function: when included, also defines #handle_return_node (instance visibility: private)

Extract the return type from an explicit ‘:return` node.

Parameters:

  • node (Parser::AST::Node) —

    the ‘:return` AST node

  • opts (Hash) —

    additional keyword options forwarded to type inference

Returns:

  • (String, nil) 


429
430
431
 
# File 'lib/docscribe/infer/returns.rb', line 429

def handle_return_node(node, **opts)
  Literals.type_from_literal(node.children.first, fallback_type: opts[:fallback_type])
end

.handle_send_node(node, **opts) ⇒ Object

Note:

module_function: when included, also defines #handle_send_node (instance visibility: private)

Handle ‘:send` node for last_expr_type.

Parameters:

  • node (Object)
  • opts (Hash)

Returns:

  • (Object) 


296
297
298
299
300
301
302
303
304
305
306
307
 
# File 'lib/docscribe/infer/returns.rb', line 296

def handle_send_node(node, **opts)
  recv = node.children[0]
  meth = node.children[1]

  if opts[:core_rbs_provider]
    rbs_type = resolve_rbs_for_send(recv, meth, opts[:core_rbs_provider], opts[:local_var_types],
                                    opts[:param_types])
    return rbs_type if rbs_type
  end

  Literals.type_from_literal(node, fallback_type: opts[:fallback_type])
end

.infer_normal_return_type(body, **opts) ⇒ String

Note:

module_function: when included, also defines #infer_normal_return_type (instance visibility: private)

Infer the normal (non-rescue) return type from a method body node.

Parameters:

  • body (Parser::AST::Node) —

    the method body AST node

  • opts (Hash) —

    additional keyword options forwarded to type inference

Returns:

  • (String) 


130
131
132
 
# File 'lib/docscribe/infer/returns.rb', line 130

def infer_normal_return_type(body, **opts)
  run_last_expr_type(body, **opts) || FALLBACK_TYPE
end

.infer_return_type(method_source) ⇒ String

Note:

module_function: when included, also defines #infer_return_type (instance visibility: private)

Infer a return type from a full method definition source string.

The source must parse to a ‘:def` or `:defs` node. If parsing fails or inference is uncertain, the fallback type is returned.

Parameters:

  • method_source (String, nil) —

    full method definition source

Returns:

  • (String)

Raises:

  • (Parser::SyntaxError) 


27
28
29
30
31
32
33
34
35
36
37
38
39
 
# File 'lib/docscribe/infer/returns.rb', line 27

def infer_return_type(method_source)
  return FALLBACK_TYPE if method_source.nil? || method_source.strip.empty?

  root = parse_method_source(method_source)
  return FALLBACK_TYPE unless root && %i[def defs].include?(root.type)

  body = root.children.last
  local_var_types = build_local_variable_types(body)
  run_last_expr_type(body, fallback_type: FALLBACK_TYPE, nil_as_optional: true,
                           local_var_types: local_var_types) || FALLBACK_TYPE
rescue Parser::SyntaxError
  FALLBACK_TYPE
end

.infer_return_type_from_node(node) ⇒ String

Note:

module_function: when included, also defines #infer_return_type_from_node (instance visibility: private)

Infer a method’s normal return type from an already parsed def/defs node.

Parameters:

  • node (Parser::AST::Node) —

    ‘:def` or `:defs` node

Returns:

  • (String) 


57
58
59
60
61
62
63
64
 
# File 'lib/docscribe/infer/returns.rb', line 57

def infer_return_type_from_node(node)
  body = extract_def_body(node)
  return FALLBACK_TYPE unless body

  local_var_types = build_local_variable_types(body)
  run_last_expr_type(body, fallback_type: FALLBACK_TYPE, nil_as_optional: true,
                           local_var_types: local_var_types) || FALLBACK_TYPE
end

.last_expr_type(node, **opts) ⇒ String?

Note:

module_function: when included, also defines #last_expr_type (instance visibility: private)

Infer the type of the last expression in a node.

Supports:

  • ‘begin` groups

  • ‘if` branches

  • ‘case` expressions

  • explicit ‘return`

  • literal-like expressions via Literals.type_from_literal

  • method calls with RBS core type lookup

Parameters:

  • node (Parser::AST::Node, nil) —

    expression node

  • fallback_type (String) —

    type used when inference is uncertain

  • nil_as_optional (Boolean) —

    whether ‘nil` unions should be rendered as optional types

  • core_rbs_provider (Object, nil) —

    optional RBS provider for core type lookup

  • param_types (Hash, nil) —

    parameter name -> type map for lvar resolution

  • local_var_types (nil) —

    pre-built local variable types map

  • opts (Hash) —

    additional keyword options forwarded to type inference

Returns:

  • (String, nil) 


402
403
404
 
# File 'lib/docscribe/infer/returns.rb', line 402

def last_expr_type(node, **opts)
  run_last_expr_type(node, **opts)
end

.lookup_lvar_type(lvar_name, local_var_types, param_types) ⇒ String?

Note:

module_function: when included, also defines #lookup_lvar_type (instance visibility: private)

Look up a local variable’s inferred type from local or parameter type maps.

Parameters:

  • lvar_name (Symbol) —

    the local variable name

  • local_var_types (Hash, nil) —

    inferred local variable type map

  • param_types (Hash, nil) —

    parameter name to type map

Returns:

  • (String, nil) 


356
357
358
359
360
361
 
# File 'lib/docscribe/infer/returns.rb', line 356

def lookup_lvar_type(lvar_name, local_var_types, param_types)
  return local_var_types[lvar_name.to_s] if local_var_types&.key?(lvar_name.to_s)
  return param_types[lvar_name.to_s] if param_types&.key?(lvar_name.to_s)

  nil
end

.parse_method_source(method_source) ⇒ Parser::AST::Node?

Note:

module_function: when included, also defines #parse_method_source (instance visibility: private)

Parse a Ruby source string into an AST using the Parser gem.

Parameters:

  • method_source (String) —

    the method definition source string to parse

Returns:

  • (Parser::AST::Node, nil) 


46
47
48
49
50
 
# File 'lib/docscribe/infer/returns.rb', line 46

def parse_method_source(method_source)
  buffer = Parser::Source::Buffer.new('(method)')
  buffer.source = method_source
  Docscribe::Parsing.parse_buffer(buffer)
end

.populate_returns_spec(spec, body, local_var_types, **opts) ⇒ Hash

Note:

module_function: when included, also defines #populate_returns_spec (instance visibility: private)

Populate the spec hash with normal and/or rescue return types from the body.

Parameters:

  • spec (Hash) —

    the return spec hash to populate

  • body (Parser::AST::Node) —

    the method body AST node

  • local_var_types (Hash, nil) —

    inferred local variable type map

  • opts (Hash) —

    additional keyword options forwarded to type inference

Returns:

  • (Hash) 


116
117
118
119
120
121
122
 
# File 'lib/docscribe/infer/returns.rb', line 116

def populate_returns_spec(spec, body, local_var_types, **opts)
  if body.type == :rescue
    process_rescue_body(spec, body, **opts)
  else
    spec[:normal] = infer_normal_return_type(body, **opts, local_var_types: local_var_types)
  end
end

.process_case_branches(node, **opts) ⇒ Array<String>

Note:

module_function: when included, also defines #process_case_branches (instance visibility: private)

Extract inferred return types from all branches of a :case expression.

Parameters:

  • node (Parser::AST::Node) —

    the :case AST node

  • opts (Hash) —

    additional keyword options forwarded to type inference

Returns:

  • (Array<String>) —

    list of inferred types from each branch



261
262
263
264
265
266
267
268
269
 
# File 'lib/docscribe/infer/returns.rb', line 261

def process_case_branches(node, **opts)
  (node.children[1..] || []).compact.flat_map do |child|
    if child.type == :when
      run_last_expr_type(child.children.last, **opts)
    else
      run_last_expr_type(child, **opts)
    end
  end.compact
end

.process_rescue_body(spec, body, **opts) ⇒ Hash

Note:

module_function: when included, also defines #process_rescue_body (instance visibility: private)

Process a :rescue body node and populate spec with normal + rescue return types.

Parameters:

  • spec (Hash) —

    the return spec hash to populate

  • body (Parser::AST::Node) —

    the :rescue AST node

  • fallback_type (String) —

    type used when inference is uncertain

  • nil_as_optional (Boolean) —

    whether nil unions render as optional types

  • core_rbs_provider (Object, nil) —

    optional RBS provider for core type lookup

  • param_types (Hash, nil) —

    parameter name to type map

  • opts (Hash) —

    additional keyword options forwarded to type inference

Returns:

  • (Hash) 


145
146
147
148
149
150
151
 
# File 'lib/docscribe/infer/returns.rb', line 145

def process_rescue_body(spec, body, **opts)
  main_body = body.children[0]
  local_var_types = build_local_variable_types(body)
  rescue_opts = opts.merge(local_var_types: local_var_types)
  spec[:normal] = run_last_expr_type(main_body, **rescue_opts) || FALLBACK_TYPE
  process_rescue_branches(spec, body, **rescue_opts)
end

.process_rescue_branches(spec, body, **opts) ⇒ Array

Note:

module_function: when included, also defines #process_rescue_branches (instance visibility: private)

Extract return types from each :resbody child and append to spec.

Parameters:

  • spec (Hash) —

    the return spec hash to populate

  • body (Parser::AST::Node) —

    the :rescue AST node

  • opts (Hash) —

    additional keyword options forwarded to type inference

Returns:

  • (Array) —

    the list of rescue type entries



160
161
162
163
164
165
166
167
168
169
 
# File 'lib/docscribe/infer/returns.rb', line 160

def process_rescue_branches(spec, body, **opts)
  body.children.each do |ch|
    next unless ch.is_a?(Parser::AST::Node) && ch.type == :resbody

    exc_list, _asgn, rescue_body = *ch
    exc_names = Raises.exception_names_from_rescue_list(exc_list)
    rtype = run_last_expr_type(rescue_body, **opts) || opts[:fallback_type]
    spec[:rescues] << [exc_names, rtype]
  end
end

.resolve_chained_send_rbs(recv, meth, core_rbs_provider, local_var_types, param_types) ⇒ String?

Note:

module_function: when included, also defines # (instance visibility: private)

Resolve RBS return type for a chained ‘:send` receiver.

Parameters:

  • recv (Object)
  • meth (Object)
  • core_rbs_provider (Object)
  • local_var_types (Object)
  • param_types (Object)

Returns:

  • (String, nil) 


373
374
375
376
377
378
379
380
381
 
# File 'lib/docscribe/infer/returns.rb', line 373

def resolve_chained_send_rbs(recv, meth, core_rbs_provider, local_var_types, param_types)
  inner_type = run_last_expr_type(recv, fallback_type: nil, nil_as_optional: false,
                                        core_rbs_provider: core_rbs_provider, param_types: param_types,
                                        local_var_types: local_var_types)
  return nil unless inner_type

  rbs_type = resolve_rbs_return_type(inner_type, meth, core_rbs_provider)
  rbs_type unless rbs_type == FALLBACK_TYPE
end

.resolve_lvar_rbs(recv, meth, core_rbs_provider, local_var_types, param_types) ⇒ String?

Note:

module_function: when included, also defines # (instance visibility: private)

Resolve RBS return type for an ‘:lvar` receiver.

Parameters:

  • recv (Object)
  • meth (Object)
  • core_rbs_provider (Object)
  • local_var_types (Object)
  • param_types (Object)

Returns:

  • (String, nil) 


340
341
342
343
344
345
346
347
 
# File 'lib/docscribe/infer/returns.rb', line 340

def resolve_lvar_rbs(recv, meth, core_rbs_provider, local_var_types, param_types)
  lvar_name = recv&.children&.first
  recv_type = lookup_lvar_type(lvar_name, local_var_types, param_types)
  return nil unless recv_type

  rbs_type = resolve_rbs_return_type(recv_type, meth, core_rbs_provider)
  rbs_type unless rbs_type == FALLBACK_TYPE
end

.resolve_rbs_for_send(recv, meth, core_rbs_provider, local_var_types, param_types) ⇒ String?

Note:

module_function: when included, also defines #resolve_rbs_for_send (instance visibility: private)

Resolve RBS return type for a send node’s receiver, if possible.

Handles ‘:lvar` and chained `:send` receivers.

Parameters:

  • recv (Parser::AST::Node, nil) —

    the receiver node of the send

  • meth (Symbol) —

    the method name being called

  • core_rbs_provider (Object, nil) —

    optional RBS provider for core type lookup

  • local_var_types (Hash, nil) —

    inferred local variable type map

  • param_types (Hash, nil) —

    parameter name to type map

Returns:

  • (String, nil) —

    resolved type or nil if unresolvable



320
321
322
323
324
325
326
327
328
 
# File 'lib/docscribe/infer/returns.rb', line 320

def resolve_rbs_for_send(recv, meth, core_rbs_provider, local_var_types, param_types)
  return nil unless core_rbs_provider

  if recv&.type == :lvar
    resolve_lvar_rbs(recv, meth, core_rbs_provider, local_var_types, param_types)
  elsif recv&.type == :send
    resolve_chained_send_rbs(recv, meth, core_rbs_provider, local_var_types, param_types)
  end
end

.resolve_rbs_return_type(container_type, method_name, core_rbs_provider) ⇒ String

Note:

module_function: when included, also defines #resolve_rbs_return_type (instance visibility: private)

Resolve an RBS return type for a method call.

Parameters:

  • container_type (String) —

    class or module name

  • method_name (String) —

    method name

  • core_rbs_provider (Object) —

    core RBS type lookup provider

Returns:

  • (String) —

    inferred return type



440
441
442
443
444
445
446
447
448
449
450
 
# File 'lib/docscribe/infer/returns.rb', line 440

def resolve_rbs_return_type(container_type, method_name, core_rbs_provider)
  return FALLBACK_TYPE unless core_rbs_provider

  sig = core_rbs_provider.signature_for(
    container: container_type,
    scope: :instance,
    name: method_name
  )

  sig&.return_type || FALLBACK_TYPE
end

.returns_spec_from_node(node, fallback_type: FALLBACK_TYPE, nil_as_optional: true, core_rbs_provider: nil, param_types: nil) ⇒ Hash

Note:

module_function: when included, also defines #returns_spec_from_node (instance visibility: private)

Return a structured return-type spec for a method node.

The result includes:

  • ‘:normal` => normal/happy-path return type

  • ‘:rescues` => array of `[exception_names, return_type]` pairs for rescue branches

Parameters:

  • node (Parser::AST::Node) —

    ‘:def` or `:defs` node

  • fallback_type (String) (defaults to: FALLBACK_TYPE) —

    type used when inference is uncertain

  • nil_as_optional (Boolean) (defaults to: true) —

    whether ‘nil` unions should be rendered as optional types

  • core_rbs_provider (nil) (defaults to: nil) —

    core RBS type lookup provider

  • param_types (nil) (defaults to: nil) —

    parameter name -> type map

Returns:

  • (Hash) 


79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
 
# File 'lib/docscribe/infer/returns.rb', line 79

def returns_spec_from_node(node, fallback_type: FALLBACK_TYPE, nil_as_optional: true, core_rbs_provider: nil,
                           param_types: nil)
  body = extract_def_body(node)
  spec = { normal: FALLBACK_TYPE, rescues: [] } #: Hash[Symbol, untyped]
  return spec unless body

  local_var_types = build_local_variable_types(body)

  populate_returns_spec(spec, body, local_var_types,
                        fallback_type: fallback_type,
                        nil_as_optional: nil_as_optional,
                        core_rbs_provider: core_rbs_provider,
                        param_types: param_types)

  spec
end

.run_last_expr_type(node, **opts) ⇒ String?

Note:

module_function: when included, also defines #run_last_expr_type (instance visibility: private)

Dispatch ‘last_expr_type` based on node type.

Parameters:

  • node (Parser::AST::Node, nil)
  • opts (Hash) —

    options passed through as keyword args

Returns:

  • (String, nil) 


412
413
414
415
416
417
418
419
420
421
 
# File 'lib/docscribe/infer/returns.rb', line 412

def run_last_expr_type(node, **opts)
  return unless node

  handler = LAST_EXPR_TYPE_HANDLERS[node.type]
  if handler
    send(handler, node, **opts)
  else
    Literals.type_from_literal(node, fallback_type: opts[:fallback_type])
  end
end

.unify_nil_types(type_a, type_b, fallback_type:, nil_as_optional:) ⇒ String

Note:

module_function: when included, also defines #unify_nil_types (instance visibility: private)

Unify two types where one may be ‘nil`, producing optional or union type.

Parameters:

  • type_a (String) —

    first type string

  • type_b (String) —

    second type string

  • fallback_type (String) —

    type used when neither is nil

  • nil_as_optional (Boolean) —

    whether to render nil unions as optional types

Returns:

  • (String) 


483
484
485
486
487
488
489
490
 
# File 'lib/docscribe/infer/returns.rb', line 483

def unify_nil_types(type_a, type_b, fallback_type:, nil_as_optional:)
  if type_a == 'nil' || type_b == 'nil'
    non_nil = (type_a == 'nil' ? type_b : type_a)
    return nil_as_optional ? "#{non_nil}?" : "#{non_nil}, nil"
  end

  fallback_type
end

.unify_types(type_a, type_b, fallback_type:, nil_as_optional:) ⇒ String?

Note:

module_function: when included, also defines #unify_types (instance visibility: private)

Unify two inferred types into a single type string.

Rules:

  • identical types remain unchanged

  • ‘nil` unions may become optional types if enabled

  • otherwise falls back conservatively to ‘fallback_type`

Parameters:

  • a (String, nil)
  • b (String, nil)
  • fallback_type (String)
  • nil_as_optional (Boolean)
  • type_a (String, nil) —

    first type to unify

  • type_b (String, nil) —

    second type to unify

Returns:

  • (String, nil) 


467
468
469
470
471
472
473
 
# File 'lib/docscribe/infer/returns.rb', line 467

def unify_types(type_a, type_b, fallback_type:, nil_as_optional:)
  type_a ||= fallback_type
  type_b ||= fallback_type
  return type_a if type_a == type_b

  unify_nil_types(type_a, type_b, fallback_type: fallback_type, nil_as_optional: nil_as_optional)
end