Module:WikitextParser

---

-- Module:WikitextParser is a general-purpose wikitext parser
-- Documentation and master version: https://en.wikipedia.org/wiki/Module:WikitextParser
-- Authors: User:Sophivorus, User:Certes, User:Aidan9382, et al.
-- License: CC-BY-SA-4.0

local str_find, str_sub, str_match = string.find, string.sub, string.match
local str_gsub, str_gmatch, str_rep = string.gsub, string.gmatch, string.rep
local str_lower = string.lower
local table_insert, table_concat = table.insert, table.concat

-- Private helper method to escape a string for use in regexes
local function escapeString( str )
    return str_gsub( str, '[%^%$%(%)%.%[%]%*%+%-%?%%]', '%%%0' )
end

-- Fast gsplit (a la mw.text.gsplit) for fixed ASCII separators
local function str_gsplit( str, sep )
    if sep == "" or str_find( sep, "[\128-\255]" ) then
        error( "Separator must be ASCII and cannot be empty" )
    end
    local i, str1 = 1, str .. sep
    return function()
        local s, e = str_find( str1, sep, i, true )
        if not s then
            return s
        end
        local t = i
        i = e + 1
        return str_sub( str1, t, s - 1 )
    end
end

-- Fast dedicated two-way split (a la mw.text.split) for fixed ASCII separators
local function str_split2( str, sep )
    if sep == "" or str_find( sep, "[\128-\255]" ) then
        error( "Separator must be ASCII and cannot be empty" )
    end
    local s, e = str_find( str, sep, 1, true )
    if not s then
        return str
    end
    return str_sub( str, 1, s - 1 ), str_sub( str, e + 1 )
end

-- Fast whitespace trim (a la mw.text.trim)
local function str_trim( str )
    return str_match( str, '^%s*(.-)%s*$' )
end

local parser = {}

-- Get the lead section from the given wikitext
-- The lead section is any content before the first section title.
-- @param wikitext Required. Wikitext to parse.
-- @return Wikitext of the lead section. May be empty if the lead section is empty.
function parser.getLead( wikitext )
    wikitext = '\n' .. wikitext
    wikitext = str_gsub( wikitext, '\n==.*', '' )
    wikitext = str_trim( wikitext )
    return wikitext
end

-- Get the sections from the given wikitext
-- This method doesn't get the lead section, use getLead for that
-- @param wikitext Required. Wikitext to parse.
-- @return Map from section title to section content
function parser.getSections( wikitext )
    local sections = {}
    wikitext = '\n' .. wikitext .. '\n=='
    for title in str_gmatch( wikitext, '\n==+ *([^=]-) *==+' ) do
        local section = str_match( wikitext, '\n==+ *' .. escapeString( title ) .. ' *==+(.-)\n==' )
        section = str_trim( section )
        sections[ title ] = section
    end
    return sections
end

-- Get a section from the given wikitext (including any subsections)
-- If the given section title appears more than once, only the section of the first instance will be returned
-- @param wikitext Required. Wikitext to parse.
-- @param title Required. Title of the section
-- @return Wikitext of the section, or nil if it isn't found. May be empty if the section is empty or contains only subsections.
function parser.getSection( wikitext, title )
    title = str_trim( title )
    title = escapeString( title )
    wikitext = '\n' .. wikitext .. '\n'
    local level
    level, wikitext = str_match( wikitext, '\n(==+) *' .. title .. ' *==.-\n(.*)' )
    if wikitext then
        local nextSection = '\n==' .. str_rep( '=?', #level - 2 ) .. '[^=].*'
        wikitext = str_gsub( wikitext, nextSection, '' ) -- remove later sections at this level or higher
        wikitext = str_trim( wikitext )
        return wikitext
    end
end

-- Get the content of a <section> tag from the given wikitext.
-- We can't use getTags because unlike all other tags, both opening and closing <section> tags are self-closing.
-- @param wikitext Required. Wikitext to parse.
-- @param name Required. Name of the <section> tag
-- @return Content of the <section> tag, or nil if it isn't found. May be empty if the section tag is empty.
function parser.getSectionTag( wikitext, name )
    name = str_trim( name )
    name = escapeString( name )
    local sections = {}
    for section in str_gmatch( wikitext, '< *section +begin *= *["\']? *' .. name .. ' *["\']? */>(.-)< *section +end= *["\']? *'.. name ..' *["\']? */>' ) do
        table_insert( sections, section )
    end
    if #sections > 0 then
        return table_concat( sections )
    end
end

-- Get the lists from the given wikitext.
-- @param wikitext Required. Wikitext to parse.
-- @return Sequence of lists.
function parser.getLists( wikitext )
    local lists = {}
    wikitext = '\n' .. wikitext .. '\n\n'
    for list in str_gmatch( wikitext, '\n([*#].-)\n[^*#]' ) do
        table_insert( lists, list )
    end
    return lists
end

-- Get the paragraphs from the given wikitext.
-- @param wikitext Required. Wikitext to parse.
-- @return Sequence of paragraphs.
function parser.getParagraphs( wikitext )
    local paragraphs = {}

    -- Remove non-paragraphs
    wikitext = '\n' .. wikitext .. '\n' -- add newlines to simplfy patterns
    wikitext = str_gsub( wikitext, '%f[^\n]<!%-%-.-%-%->%f[\n]', '' ) -- remove comments
    wikitext = str_gsub( wikitext, '%f[^\n]%[%b[]%]%f[\n]', '' ) -- remove files and categories
    wikitext = str_gsub( wikitext, '%f[^\n]%b{} *%f[\n]', '' ) -- remove tables and block templates
    wikitext = str_gsub( wikitext, '%f[^\n]%b{} *%b{} *%f[\n]', '' ) -- remove neighboring tables and block templates
    wikitext = str_gsub( wikitext, '%f[^\n]%b{} *<!%-%-.-%-%-> *%b{} *%f[\n]', '' ) -- remove neighboring tables and block templates with a comment among them
    wikitext = str_gsub( wikitext, '%f[^\n][*#].-%f[\n]', '' ) -- remove lists
    wikitext = str_gsub( wikitext, '%f[^\n]==+[^=]+==+ *%f[\n]', '' ) -- remove section titles
    wikitext = str_trim( wikitext )

    for paragraph in str_gsplit( wikitext, "\n\n" ) do
        paragraph = str_trim( paragraph )
        if paragraph ~= "" then -- Filter out 3+ consecutive newlines
            table_insert( paragraphs, paragraph )
        end
    end
    return paragraphs
end

-- Get the templates from the given wikitext.
-- @param wikitext Required. Wikitext to parse.
-- @return Sequence of templates.
function parser.getTemplates( wikitext )
    local templates = {}
    for template in str_gmatch( wikitext, '{%b{}}' ) do
        if str_sub( template, 1, 3 ) ~= '{{#' then -- skip parser functions like #if
            table_insert( templates, template )
        end
    end
    return templates
end

-- Get the requested template from the given wikitext.
-- If the template appears more than once, only the first instance will be returned
-- @param wikitext Required. Wikitext to parse.
-- @param name Name of the template to get
-- @return Wikitext of the template, or nil if it wasn't found
function parser.getTemplate( wikitext, name )
    local templates = parser.getTemplates( wikitext )
    local lang = mw.language.getContentLanguage()
    for _, template in ipairs( templates ) do
        local templateName = parser.getTemplateName( template )
        if lang:ucfirst( templateName ) == lang:ucfirst( name ) then
            return template
        end
    end
end

-- Get name of the template from the given template wikitext.
-- @param templateWikitext Required. Wikitext of the template to parse.
-- @return Name of the template
-- @todo Strip "Template:" namespace?
function parser.getTemplateName( templateWikitext )
    return str_match( templateWikitext, '^{{ *([^}|\n]+)' )
end

-- Get the parameters from the given template wikitext.
-- @param templateWikitext Required. Wikitext of the template to parse.
-- @return Map from parameter names to parameter values, NOT IN THE ORIGINAL ORDER.
-- @return Order in which the parameters were parsed.
function parser.getTemplateParameters( templateWikitext )
    local parameters = {}
    local paramOrder = {}
    local params = str_match( templateWikitext, '{{[^|}]-|(.*)}}' )
    if params then
        -- Temporarily replace pipes in subtemplates and links to avoid chaos
        for subtemplate in str_gmatch( params, '{%b{}}' ) do
            params = str_gsub( params, escapeString( subtemplate ), str_gsub( subtemplate, '.', { ['%']='%%', ['|']="@@:@@", ['=']='@@_@@' } ) )
        end
        for link in str_gmatch( params, '%[%b[]%]' ) do
            params = str_gsub( params, escapeString( link ), str_gsub( link, '.', { ['%']='%%', ['|']='@@:@@', ['=']='@@_@@' } ) )
        end
        local count = 0
        local name, value
        for param in str_gsplit( params, '|' ) do
            name, value = str_split2( param, '=' )
            if value then
                local number = tonumber( name )
                if number then --XXX: What about keys like: 01, .1, 2e0, etc.
                    name = number
                end
                value = str_trim( value )
            else
                value = name
                count = count + 1
                name = count
            end
            value = str_gsub( value, '@@_@@', '=' )
            value = str_gsub( value, '@@:@@', '|' )
            parameters[ name ] = value
            table_insert( paramOrder, name )
        end
    end
    return parameters, paramOrder
end

-- Get the tags from the given wikitext.
-- @param wikitext Required. Wikitext to parse.
-- @return Sequence of tags.
function parser.getTags( wikitext )
    local tags = {}
    local tag, tagName, tagEnd
    -- Don't match closing tags like </div>, comments like <!--foo-->, comparisons like 1<2 or things like <3
    for tagStart, tagOpen in str_gmatch( wikitext, '()(<[^/!%d].->)' ) do
        tagName = parser.getTagName( tagOpen )

        -- If we're in a self-closing tag, like <ref name="foo" />, <references/>, <br/>, <br>, <hr>, etc.
        if str_match( tagOpen, '<.-/>' ) or tagName == 'br' or tagName == 'hr' then
            tag = tagOpen

        -- If we're in a tag that may contain others like it, like <div> or <span>
        elseif tagName == 'div' or tagName == 'span' then
            local position = tagStart + #tagOpen - 1
            local depth = 1
            while depth > 0 do
                tagEnd = str_match( wikitext, '</ ?' .. tagName .. ' ?>()', position )
                if tagEnd then
                    tagEnd = tagEnd - 1
                else
                    break -- unclosed tag
                end 
                position = str_match( wikitext, '()< ?' .. tagName .. '[ >]', position + 1 )
                if not position then
                    position = tagEnd + 1
                end
                if position > tagEnd then
                    depth = depth - 1
                else
                    depth = depth + 1
                end
            end
            tag = str_sub( wikitext, tagStart, tagEnd )

        -- Else we're probably in tag that shouldn't contain others like it, like <math> or <strong>
        else
            tagEnd = str_match( wikitext, '</ ?' .. tagName .. ' ?>()', tagStart )
            if tagEnd then
                tag = str_sub( wikitext, tagStart, tagEnd - 1 )

            -- If no end tag is found, assume we matched something that wasn't a tag, like <no. 1>
            else
                tag = nil
            end
        end
        table_insert( tags, tag )
    end
    return tags
end

-- Get the name of the tag in the given wikitext
-- @param tag Required. Tag to parse.
-- @return Name of the tag or nil if not found
function parser.getTagName( tagWikitext )
    local tagName = str_match( tagWikitext, '^< *(.-)[ />]' )
    if tagName then tagName = str_lower( tagName ) end
    return tagName
end

-- Get the value of an attribute in the given tag.
-- @param tagWikitext Required. Wikitext of the tag to parse.
-- @param attribute Required. Name of the attribute.
-- @return Value of the attribute or nil if not found
function parser.getTagAttribute( tagWikitext, attribute )
    local _, value = str_match( tagWikitext, '^<[^/>]*' .. attribute .. ' *= *(["\']?)([^/>]-)%1[ />]' )
    return value
end

-- Get the content of the given tag.
-- @param tagWikitext Required. Wikitext of the tag to parse.
-- @return Content of the tag. May be empty if the tag is empty. Will be nil if the tag is self-closing.
-- @todo May fail with nested tags
function parser.getTagContent( tagWikitext )
    return str_match( tagWikitext, '^<.->.-</.->' )
end

-- Get the <gallery> tags from the given wikitext.
-- @param wikitext Required. Wikitext to parse.
-- @return Sequence of gallery tags.
function parser.getGalleries( wikitext )
    local galleries = {}
    local tags = parser.getTags( wikitext )
    for _, tag in ipairs( tags ) do
        local tagName = parser.getTagName( tag )
        if tagName == 'gallery' then
            table_insert( galleries, tag )
        end
    end
    return galleries
end

-- Get the <ref> tags from the given wikitext.
-- @param wikitext Required. Wikitext to parse.
-- @return Sequence of ref tags.
function parser.getReferences( wikitext )
    local references = {}
    local tags = parser.getTags( wikitext )
    for _, tag in ipairs( tags ) do
        local tagName = parser.getTagName( tag )
        if tagName == 'ref' then
            table_insert( references, tag )
        end
    end
    return references
end

-- Get the reference with the given name from the given wikitext.
-- @param wikitext Required. Wikitext to parse.
-- @param referenceName Required. Name of the reference.
-- @return Wikitext of the reference
function parser.getReference( wikitext, referenceName )
    local references = parser.getReferences( wikitext )
    for _, reference in ipairs( references ) do
        local content = parser.getTagContent( reference )
        local name = parser.getTagAttribute( reference, 'name' )
        if content and name == referenceName then
            return reference
        end
    end
end

-- Get the tables from the given wikitext.
-- @param wikitext Required. Wikitext to parse.
-- @return Sequence of tables.
function parser.getTables( wikitext )
    local tables = {}
    wikitext = '\n' .. wikitext
    for t in str_gmatch( wikitext, '\n%b{}' ) do
        if str_sub( t, 1, 3 ) == '\n{|' then
            t = str_trim( t ) -- exclude the leading newline
            table_insert( tables, t )
        end
    end
    return tables
end

-- Get the id from the given table wikitext
-- @param tableWikitext Required. Wikitext of the table to parse.
-- @param attribute Required. Name of the attribute.
-- @return Value of the attribute or nil if not found
function parser.getTableAttribute( tableWikitext, attribute )
    local _, value = str_match( tableWikitext, '^{|[^\n]*' .. attribute .. ' *= *(["\']?)([^\n]-)%1[^\n]*\n' )
    if not value or value == '' then
        value = str_match( tableWikitext, '^{|[^\n]*' .. attribute .. ' *= *([^\n ]+)[^\n]*\n' )
    end
    return value
end

-- Get a table by id from the given wikitext
-- @param wikitext Required. Wikitext to parse.
-- @param id Required. Id of the table
-- @return Wikitext of the table or nil if not found
function parser.getTable( wikitext, id )
    local tables = parser.getTables( wikitext )
    for _, t in ipairs( tables ) do
        if id == parser.getTableAttribute( t, 'id' ) then
            return t
        end
    end
end

-- Get the data from the given table wikitext
-- @param tableWikitext Required. Wikitext of the table to parse.
-- @return Table data
-- @todo Test and make more robust
function parser.getTableData( tableWikitext )
    local tableData = {}
    tableWikitext = str_trim( tableWikitext );
    tableWikitext = str_gsub( tableWikitext, '^{|.-\n', '' ) -- remove the header
    tableWikitext = str_gsub( tableWikitext, '\n|}$', '' ) -- remove the footer
    tableWikitext = str_gsub( tableWikitext, '^|%+.-\n', '' ) -- remove any caption
    tableWikitext = str_gsub( tableWikitext, '|%-.-\n', '|-\n' ) -- remove any row attributes
    tableWikitext = str_gsub( tableWikitext, '^|%-\n', '' ) -- remove any leading empty row
    tableWikitext = str_gsub( tableWikitext, '\n|%-$', '' ) -- remove any trailing empty row
    for rowWikitext in str_gsplit( tableWikitext, '|-' ) do
        local rowData = {}
        rowWikitext = str_gsub( rowWikitext, '||', '\n|' )
        rowWikitext = str_gsub( rowWikitext, '!!', '\n|' )
        rowWikitext = str_gsub( rowWikitext, '\n!', '\n|' )
        rowWikitext = str_gsub( rowWikitext, '^!', '\n|' )
        rowWikitext = str_gsub( rowWikitext, '^\n|', '' )
        for cellWikitext in str_gsplit( rowWikitext, '\n|' ) do
            cellWikitext = str_trim( cellWikitext )
            table_insert( rowData, cellWikitext )
        end
        table_insert( tableData, rowData )
    end
    return tableData
end

-- Get the internal links from the given wikitext (includes category and file links).
-- @param wikitext Required. Wikitext to parse.
-- @return Sequence of internal links.
function parser.getLinks( wikitext )
    local links = {}
    for link in str_gmatch( wikitext, '%[%b[]%]' ) do
        table_insert( links, link )
    end
    return links
end

-- Get the file links from the given wikitext.
-- @param wikitext Required. Wikitext to parse.
-- @return Sequence of file links.
function parser.getFiles( wikitext )
    local files = {}
    local links = parser.getLinks( wikitext )
    for _, link in ipairs( links ) do
        local namespace = str_match( link, '^%[%[ *(.-) *:' )
        if namespace and mw.site.namespaces[ namespace ] and mw.site.namespaces[ namespace ].canonicalName == 'File' then
            table_insert( files, link )
        end
    end
    return files
end

-- Get name of the file from the given file wikitext.
-- @param fileWikitext Required. Wikitext of the file to parse.
-- @return Name of the file
function parser.getFileName( fileWikitext )
    return str_match( fileWikitext, '^%[%[ *.- *: *(.-) *[]|]' )
end

-- Get the category links from the given wikitext.
-- @param wikitext Required. Wikitext to parse.
-- @return Sequence of category links.
function parser.getCategories( wikitext )
    local categories = {}
    local links = parser.getLinks( wikitext )
    for _, link in ipairs( links ) do
        local namespace = str_match( link, '^%[%[ -(.-) -:' )
        if namespace and mw.site.namespaces[ namespace ] and mw.site.namespaces[ namespace ].canonicalName == 'Category' then
            table_insert( categories, link )
        end
    end
    return categories
end

-- Get the external links from the given wikitext.
-- @param wikitext Required. Wikitext to parse.
-- @return Sequence of external links.
function parser.getExternalLinks( wikitext )
    local links = {}
    for link in str_gmatch( wikitext, '%b[]' ) do
        if str_match( link, '^%[//' ) or str_match( link, '^%[https?://' ) then
            table_insert( links, link )
        end
    end
    return links
end

return parser