diff --git a/src/parse/parse_parameters.ts b/src/parse/parse_parameters.ts index a2c6122..2636acc 100644 --- a/src/parse/parse_parameters.ts +++ b/src/parse/parse_parameters.ts @@ -115,7 +115,7 @@ function parseYields(parameters: string[], body: string[]): Yields { } function parseReturnFromDefinition(parameters: string[]): Returns | null { - const pattern = /^->\s*([\w\[\], \.]*)/; + const pattern = /^->\s*(["']?)([\w\[\], \.]*)\1/; for (const param of parameters) { const match = param.trim().match(pattern); @@ -125,7 +125,7 @@ function parseReturnFromDefinition(parameters: string[]): Returns | null { } // Skip "-> None" annotations - return match[1] === "None" ? null : { type: match[1] }; + return match[1] === "None" ? null : { type: match[2] }; } return null; diff --git a/src/parse/tokenize_definition.ts b/src/parse/tokenize_definition.ts index 2c1b7dc..c170143 100644 --- a/src/parse/tokenize_definition.ts +++ b/src/parse/tokenize_definition.ts @@ -1,5 +1,5 @@ export function tokenizeDefinition(functionDefinition: string): string[] { - const definitionPattern = /(?:def|class)\s+\w+\s*\(([\s\S]*)\)\s*(->\s*[\w\[\], \.]*)?:\s*(?:#.*)?$/; + const definitionPattern = /(?:def|class)\s+\w+\s*\(([\s\S]*)\)\s*(->\s*(["']?)[\w\[\], \.]*\3)?:\s*(?:#.*)?$/; const match = definitionPattern.exec(functionDefinition); if (match == undefined || match[1] == undefined) { diff --git a/src/parse/utilities.ts b/src/parse/utilities.ts index c99317b..a223384 100644 --- a/src/parse/utilities.ts +++ b/src/parse/utilities.ts @@ -15,8 +15,8 @@ export function getIndentation(line: string): string { */ export function preprocessLines(lines: string[]): string[] { return lines - .map(line => line.trim()) - .filter((line) => !line.startsWith("#")); + .filter((line) => !line.trim().startsWith("#")) + .map(line => line.split("#")[0].trim()); } export function indentationOf(line: string): number { diff --git a/src/test/parse/utilities.spec.ts b/src/test/parse/utilities.spec.ts index 79043bd..48542c1 100644 --- a/src/test/parse/utilities.spec.ts +++ b/src/test/parse/utilities.spec.ts @@ -57,4 +57,19 @@ describe("preprocessLines()", () => { 'foo' ]); }); + + it('should discard inline comments.', () => { + const result = preprocessLines([ + 'foo', + 'arg: int # hello', + 'hello#world', + 'hello\t\t#world' + ]); + expect(result).to.be.deep.equal([ + 'foo', + 'arg: int', + 'hello', + 'hello' + ]); + }) });